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

use std::io::Result as IoResult;
use std::io::Write;

/// Splits the incoming data into HTTP chunks.
///
/// # Example
///
/// ```
/// use chunked_transfer::Encoder;
/// use std::io::Write;
///
/// let mut decoded = "hello world";
/// let mut encoded: Vec<u8> = vec![];
///
/// {
///     let mut encoder = Encoder::with_chunks_size(&mut encoded, 5);
///     encoder.write_all(decoded.as_bytes());
/// }
///
/// assert_eq!(encoded, b"5\r\nhello\r\n5\r\n worl\r\n1\r\nd\r\n0\r\n\r\n");
/// ```
pub struct Encoder<W>
where
    W: Write,
{
    // where to send the result
    output: W,

    // size of each chunk
    chunks_size: usize,

    // data waiting to be sent is stored here
    // This will always be at least 6 bytes long. The first 6 bytes
    // are reserved for the chunk size and \r\n.
    buffer: Vec<u8>,

    // Flushes the internal buffer after each write. This might be useful
    // if data should be sent immediately to downstream consumers
    flush_after_write: bool,
}

const MAX_CHUNK_SIZE: usize = std::u32::MAX as usize;
const HEADER_SIZE: usize = 6;

impl<W> Encoder<W>
where
    W: Write,
{
    pub fn new(output: W) -> Encoder<W> {
        Encoder::with_chunks_size(output, 8192)
    }

    pub fn with_chunks_size(output: W, chunks: usize) -> Encoder<W> {
        let chunks_size = if chunks > MAX_CHUNK_SIZE {
            MAX_CHUNK_SIZE
        } else {
            chunks
        };
        let mut encoder = Encoder {
            output,
            chunks_size,
            buffer: Vec::with_capacity(0),
            flush_after_write: false,
        };
        encoder.reset_buffer();
        encoder
    }

    pub fn with_flush_after_write(output: W) -> Encoder<W> {
        let mut encoder = Encoder {
            output,
            chunks_size: 8192,
            buffer: Vec::with_capacity(0),
            flush_after_write: true,
        };
        encoder.reset_buffer();
        encoder
    }

    fn reset_buffer(&mut self) {
        // Reset buffer
        self.buffer.clear();
        // Arbitrary bytes to take up space. Should never be seen in output.
        let placeholder: [u8; HEADER_SIZE] = *b"DACE\r\n";
        self.buffer.extend_from_slice(&placeholder);
    }

    fn buffer_empty(&self) -> bool {
        self.buffer.len() == HEADER_SIZE
    }

    fn send(&mut self) -> IoResult<()> {
        // Never send an empty buffer, because that would be interpreted
        // as the end of the stream, which we indicate explicitly on drop.
        if self.buffer_empty() {
            return Ok(());
        }
        // Prepend the length and \r\n to the buffer.
        let prelude = format!("{:x}\r\n", self.buffer.len() - HEADER_SIZE);
        let prelude = prelude.as_bytes();
        if prelude.len() > HEADER_SIZE {
            // This should never happen because MAX_CHUNK_SIZE of u32::MAX
            // can always be encoded in 4 hex bytes.
            panic!("invariant failed: prelude longer than HEADER_SIZE");
        }
        // Copy the prelude into the buffer. For small chunks, this won't necessarily
        // take up all the space that was reserved for the prelude.
        let offset = HEADER_SIZE - prelude.len();
        self.buffer[offset..HEADER_SIZE].clone_from_slice(&prelude);

        // Append the chunk-finishing \r\n to the buffer.
        self.buffer.write_all(b"\r\n")?;

        self.output.write_all(&self.buffer[offset..])?;
        self.reset_buffer();

        Ok(())
    }
}

impl<W> Write for Encoder<W>
where
    W: Write,
{
    fn write(&mut self, data: &[u8]) -> IoResult<usize> {
        println!("writing to buffer: {:?}", data);
        let buffer_size = self.buffer.len() - HEADER_SIZE;
        let remaining_buffer_space = self.chunks_size - buffer_size;
        let bytes_to_buffer = std::cmp::min(remaining_buffer_space, data.len());
        self.buffer.extend_from_slice(&data[0..bytes_to_buffer]);
        let more_to_write: bool = bytes_to_buffer < data.len();
        if self.flush_after_write || more_to_write {
            self.send()?;
        }

        // If we didn't write the whole thing, keep working on it.
        if more_to_write {
            self.write_all(&data[bytes_to_buffer..])?;
        }
        Ok(data.len())
    }

    fn flush(&mut self) -> IoResult<()> {
        self.send()
    }
}

impl<W> Drop for Encoder<W>
where
    W: Write,
{
    fn drop(&mut self) {
        self.flush().ok();
        write!(self.output, "0\r\n\r\n").ok();
    }
}

#[cfg(test)]
mod test {
    use super::Encoder;
    use std::io;
    use std::io::Write;
    use std::str::from_utf8;

    #[test]
    fn test() {
        let mut source = io::Cursor::new("hello world".to_string().into_bytes());
        let mut dest: Vec<u8> = vec![];

        {
            let mut encoder = Encoder::with_chunks_size(dest.by_ref(), 5);
            io::copy(&mut source, &mut encoder).unwrap();
            assert!(!encoder.buffer_empty());
        }

        let output = from_utf8(&dest).unwrap();

        assert_eq!(output, "5\r\nhello\r\n5\r\n worl\r\n1\r\nd\r\n0\r\n\r\n");
    }
    #[test]
    fn flush_after_write() {
        let mut source = io::Cursor::new("hello world".to_string().into_bytes());
        let mut dest: Vec<u8> = vec![];

        {
            let mut encoder = Encoder::with_flush_after_write(dest.by_ref());
            io::copy(&mut source, &mut encoder).unwrap();
            // The internal buffer has been flushed.
            assert!(encoder.buffer_empty());
        }

        let output = from_utf8(&dest).unwrap();

        assert_eq!(output, "b\r\nhello world\r\n0\r\n\r\n");
    }
}
<commit_msg>Use chunks.min(MAX_CHUNK_SIZE)<commit_after>// Copyright 2015 The tiny-http Contributors
// Copyright 2015 The rust-chunked-transfer Contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//	http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::io::Result as IoResult;
use std::io::Write;

/// Splits the incoming data into HTTP chunks.
///
/// # Example
///
/// ```
/// use chunked_transfer::Encoder;
/// use std::io::Write;
///
/// let mut decoded = "hello world";
/// let mut encoded: Vec<u8> = vec![];
///
/// {
///     let mut encoder = Encoder::with_chunks_size(&mut encoded, 5);
///     encoder.write_all(decoded.as_bytes());
/// }
///
/// assert_eq!(encoded, b"5\r\nhello\r\n5\r\n worl\r\n1\r\nd\r\n0\r\n\r\n");
/// ```
pub struct Encoder<W>
where
    W: Write,
{
    // where to send the result
    output: W,

    // size of each chunk
    chunks_size: usize,

    // data waiting to be sent is stored here
    // This will always be at least 6 bytes long. The first 6 bytes
    // are reserved for the chunk size and \r\n.
    buffer: Vec<u8>,

    // Flushes the internal buffer after each write. This might be useful
    // if data should be sent immediately to downstream consumers
    flush_after_write: bool,
}

const MAX_CHUNK_SIZE: usize = std::u32::MAX as usize;
const HEADER_SIZE: usize = 6;

impl<W> Encoder<W>
where
    W: Write,
{
    pub fn new(output: W) -> Encoder<W> {
        Encoder::with_chunks_size(output, 8192)
    }

    pub fn with_chunks_size(output: W, chunks: usize) -> Encoder<W> {
        let chunks_size = chunks.min(MAX_CHUNK_SIZE);
        let mut encoder = Encoder {
            output,
            chunks_size,
            buffer: Vec::with_capacity(0),
            flush_after_write: false,
        };
        encoder.reset_buffer();
        encoder
    }

    pub fn with_flush_after_write(output: W) -> Encoder<W> {
        let mut encoder = Encoder {
            output,
            chunks_size: 8192,
            buffer: Vec::with_capacity(0),
            flush_after_write: true,
        };
        encoder.reset_buffer();
        encoder
    }

    fn reset_buffer(&mut self) {
        // Reset buffer
        self.buffer.clear();
        // Arbitrary bytes to take up space. Should never be seen in output.
        let placeholder: [u8; HEADER_SIZE] = *b"DACE\r\n";
        self.buffer.extend_from_slice(&placeholder);
    }

    fn buffer_empty(&self) -> bool {
        self.buffer.len() == HEADER_SIZE
    }

    fn send(&mut self) -> IoResult<()> {
        // Never send an empty buffer, because that would be interpreted
        // as the end of the stream, which we indicate explicitly on drop.
        if self.buffer_empty() {
            return Ok(());
        }
        // Prepend the length and \r\n to the buffer.
        let prelude = format!("{:x}\r\n", self.buffer.len() - HEADER_SIZE);
        let prelude = prelude.as_bytes();
        if prelude.len() > HEADER_SIZE {
            // This should never happen because MAX_CHUNK_SIZE of u32::MAX
            // can always be encoded in 4 hex bytes.
            panic!("invariant failed: prelude longer than HEADER_SIZE");
        }
        // Copy the prelude into the buffer. For small chunks, this won't necessarily
        // take up all the space that was reserved for the prelude.
        let offset = HEADER_SIZE - prelude.len();
        self.buffer[offset..HEADER_SIZE].clone_from_slice(&prelude);

        // Append the chunk-finishing \r\n to the buffer.
        self.buffer.write_all(b"\r\n")?;

        self.output.write_all(&self.buffer[offset..])?;
        self.reset_buffer();

        Ok(())
    }
}

impl<W> Write for Encoder<W>
where
    W: Write,
{
    fn write(&mut self, data: &[u8]) -> IoResult<usize> {
        println!("writing to buffer: {:?}", data);
        let buffer_size = self.buffer.len() - HEADER_SIZE;
        let remaining_buffer_space = self.chunks_size - buffer_size;
        let bytes_to_buffer = std::cmp::min(remaining_buffer_space, data.len());
        self.buffer.extend_from_slice(&data[0..bytes_to_buffer]);
        let more_to_write: bool = bytes_to_buffer < data.len();
        if self.flush_after_write || more_to_write {
            self.send()?;
        }

        // If we didn't write the whole thing, keep working on it.
        if more_to_write {
            self.write_all(&data[bytes_to_buffer..])?;
        }
        Ok(data.len())
    }

    fn flush(&mut self) -> IoResult<()> {
        self.send()
    }
}

impl<W> Drop for Encoder<W>
where
    W: Write,
{
    fn drop(&mut self) {
        self.flush().ok();
        write!(self.output, "0\r\n\r\n").ok();
    }
}

#[cfg(test)]
mod test {
    use super::Encoder;
    use std::io;
    use std::io::Write;
    use std::str::from_utf8;

    #[test]
    fn test() {
        let mut source = io::Cursor::new("hello world".to_string().into_bytes());
        let mut dest: Vec<u8> = vec![];

        {
            let mut encoder = Encoder::with_chunks_size(dest.by_ref(), 5);
            io::copy(&mut source, &mut encoder).unwrap();
            assert!(!encoder.buffer_empty());
        }

        let output = from_utf8(&dest).unwrap();

        assert_eq!(output, "5\r\nhello\r\n5\r\n worl\r\n1\r\nd\r\n0\r\n\r\n");
    }
    #[test]
    fn flush_after_write() {
        let mut source = io::Cursor::new("hello world".to_string().into_bytes());
        let mut dest: Vec<u8> = vec![];

        {
            let mut encoder = Encoder::with_flush_after_write(dest.by_ref());
            io::copy(&mut source, &mut encoder).unwrap();
            // The internal buffer has been flushed.
            assert!(encoder.buffer_empty());
        }

        let output = from_utf8(&dest).unwrap();

        assert_eq!(output, "b\r\nhello world\r\n0\r\n\r\n");
    }
}
<|endoftext|>


<commit_before>use memory::{EmptyMemory, Memory, pack_u16, high_byte, low_byte, low_nibble, high_nibble};
use extensions::Incrementor;
use cpu::{Register, Flags, CarryFlag, HalfCarryFlag, ZeroFlag, SubtractFlag};
use ops::u8_immediate;

fn half_carry_for_subtract(val1: u8, val2: u8, carry: u8) -> bool {
    return (low_nibble(val2) + carry) > low_nibble(val1);
}

pub fn sub_u8_immediate(mem: &Memory, pc: &mut Register<u16>, reg: &mut Register<u8>, freg: &mut Register<Flags>, with_carry: bool) {
    let val = u8_immediate(mem, pc);
    internal_sub(reg, val, freg, with_carry);
}

fn carry_for_subtract(val1: u8, val2: u8, carry: u8) -> bool {
    return (val2 as u16 + carry as u16) > val1 as u16 ;
}

pub fn internal_sub(reg: &mut Register<u8>, val: u8, freg: &mut Register<Flags>, with_carry: bool) {
    let reg_val = reg.read();
    let mut flags = SubtractFlag;
    let mut carry = 0;

    if freg.read().contains(CarryFlag) && with_carry {
        carry = 1;
    }

    if half_carry_for_subtract(reg_val, val, carry) {
        flags.insert(HalfCarryFlag);
    }

    if carry_for_subtract(reg_val, val, carry) {
        flags.insert(CarryFlag);
    }

    let result = reg_val.wrapping_sub(val).wrapping_sub(carry);

    if result == 0 {
        flags.insert(ZeroFlag);
    }
    debug!("internal sub result {:X}", result);
    reg.write(result);
    freg.write(flags);
}

pub fn sub(reg: &mut Register<u8>, val: u8, freg: &mut Register<Flags>) {
    internal_sub(reg, val, freg, false);
}

pub fn sbc(reg: &mut Register<u8>, val: u8, freg: &mut Register<Flags>) {
    internal_sub(reg, val, freg, true);
}

pub fn sub_value_at_address(mem: &Memory, reg: &mut Register<u8>, addr: u16, freg: &mut Register<Flags>) {
    let val = mem.read_byte(addr);
    internal_sub(reg, val, freg, false);
}

pub fn sbc_value_at_address(mem: &Memory, reg: &mut Register<u8>, addr: u16, freg: &mut Register<Flags>) {
    let val = mem.read_byte(addr);
    internal_sub(reg, val, freg, true);
}

pub fn compare(reg: &mut Register<u8>, val: u8, freg: &mut Register<Flags>) {
    let reg_val = reg.read();
    let mut flags = SubtractFlag;
    let mut carry = 0;

    if half_carry_for_subtract(reg_val, val, carry) {
        flags.insert(HalfCarryFlag);
    }

    if carry_for_subtract(reg_val, val, carry) {
        flags.insert(CarryFlag);
    }

    if (reg_val == val) {
        flags.insert(ZeroFlag);
    }
    freg.write(flags);
}

pub fn compare_value_at_address(mem: &Memory, reg: &mut Register<u8>, addr: u16, freg: &mut Register<Flags>) {
    let val = mem.read_byte(addr);
    compare(reg, val, freg);
}

pub fn decrement_register(reg: &mut Register<u8>, freg: &mut Register<Flags>) {
    internal_sub(reg, 1, freg, false);
}

pub fn decrement_register_pair(reg1: &mut Register<u8>, reg2: &mut Register<u8>) {
    let val = pack_u16(reg1.read(), reg2.read());
    let ans = val - 1;

    reg1.write(high_byte(ans));
    reg2.write(low_byte(ans));
}

pub fn decrement_value_at_address(mem: &mut Memory, hb: u8, lb: u8, freg: &mut Register<Flags>) {
    let addr = pack_u16(hb, lb);
    let val = mem.read_byte(addr);
    let mut reg = Register::new(val);
    decrement_register(&mut reg, freg);
    mem.write_byte(addr, reg.read());
}

#[cfg(test)]
mod tests {
    use super::*;
    use memory::{EmptyMemory, Memory, pack_u16, high_byte, low_byte, low_nibble, high_nibble};
    use extensions::Incrementor;
    use cpu::{Register, Flags, CarryFlag, HalfCarryFlag, ZeroFlag, SubtractFlag};

    #[test]
    fn test_sub_u8_immediate() {
        let mut mem = EmptyMemory::new(0xFFFF);
        let mut pc = Register::new(0x2736);
        let mut reg = Register::new(0x35);
        let mut freg = Register::new(Flags::empty());
        mem.write_byte(pc.read(), 0x11);

        sub_u8_immediate(&mem, &mut pc, &mut reg, &mut freg, false);

        assert!(pc.read() == 0x2737);
        assert!(reg.read() == 0x24);
        assert!(freg.read() == SubtractFlag);

        mem.write_byte(pc.read(), 0x0C);
        sub_u8_immediate(&mem, &mut pc, &mut reg, &mut freg, false);

        assert!(pc.read() == 0x2738);
        assert!(reg.read() == 0x18);
        assert!(freg.read() == SubtractFlag | HalfCarryFlag);

        mem.write_byte(pc.read(), 0x0F);
        freg.write(CarryFlag);
        sub_u8_immediate(&mem, &mut pc, &mut reg, &mut freg, true);

        assert!(pc.read() == 0x2739);
        assert!(reg.read() == 0x08);
        assert!(freg.read() == SubtractFlag | HalfCarryFlag);
    }

    #[test]
    fn test_cp() {
        let mut reg = Register::new(0xAA);
        let mut val = 0xBB;
        let mut freg = Register::new(Flags::empty());

        compare(&mut reg, val, &mut freg);
        assert!(reg.read() == 0xAA);
        assert!(freg.read() == SubtractFlag | HalfCarryFlag | CarryFlag);

        reg.write(0xF0);
        val = 0xF0;
        compare(&mut reg, val, &mut freg);
        assert!(reg.read() == 0xF0);
        assert!(freg.read() == SubtractFlag | ZeroFlag);
    }

    #[test]
    fn test_compare_value_at_address() {
        let mut mem = EmptyMemory::new(0xFFFF);
        let mut reg = Register::new(0xBA);
        let mut freg = Register::new(Flags::empty());
        let addr = 0x1239;
        let mut val = 0xAB;
        mem.write_byte(addr, val);

        compare_value_at_address(&mem, &mut reg, addr, &mut freg);
        assert!(reg.read() == 0xBA);
        assert!(freg.read() == SubtractFlag | HalfCarryFlag);

        mem.write_byte(addr, 0xCD);
        reg.write(0xCD);
        compare_value_at_address(&mem, &mut reg, addr, &mut freg);
        assert!(reg.read() == 0xCD);
        assert!(freg.read() == SubtractFlag | ZeroFlag);

        mem.write_byte(addr, 0xCD);
        reg.write(0xCF);
        compare_value_at_address(&mem, &mut reg, addr, &mut freg);
        assert!(reg.read() == 0xCD);
        assert!(freg.read() == SubtractFlag | ZeroFlag | CarryFlag | HalfCarryFlag);

    }

    #[test]
    fn test_sbc_value_at_address() {
        let mut mem = EmptyMemory::new(0xFFFF);
        let mut freg = Register::new(CarryFlag);
        let mut reg = Register::new(0xFF);
        let addr = 0x1234;
        mem.write_byte(addr, 0xCD);

        sbc_value_at_address(&mem, &mut reg, addr, &mut freg);
        assert!(reg.read() == 0x31);
        assert!(freg.read() == SubtractFlag);
    }

    #[test]
    fn test_sub_value_at_address() {
        let mut mem = EmptyMemory::new(0xFFFF);
        let mut freg = Register::new(CarryFlag);
        let mut reg = Register::new(0xFF);
        let addr = 0x1234;
        mem.write_byte(addr, 0xCD);

        sub_value_at_address(&mem, &mut reg, addr, &mut freg);
        assert!(reg.read() == 0x32);
        assert!(freg.read() == SubtractFlag);
    }

    #[test]
    fn test_sbc() {
        let mut reg = Register::new(0xFF);
        let mut freg = Register::new(Flags::empty());

        sbc(&mut reg, 0x0F, &mut freg);
        assert!(reg.read() == 0xF0);
        assert!(freg.read() == SubtractFlag);

        reg.write(0xFF);
        freg.write(CarryFlag);
        sbc(&mut reg, 0x0F, &mut freg);
        assert!(reg.read() == 0xEF);
        assert!(freg.read() == SubtractFlag | HalfCarryFlag);

        reg.write(0xFF);
        freg.write(CarryFlag);
        sbc(&mut reg, 0xFF, &mut freg);
        assert!(reg.read() == 0xFF);
        assert!(freg.read() == SubtractFlag | HalfCarryFlag | CarryFlag);

        reg.write(0xAB);
        freg.write(CarryFlag);
        sbc(&mut reg, 0x12, &mut freg);
        assert!(reg.read() == 0x98);
        assert!(freg.read() == SubtractFlag);
    }

    #[test]
    fn test_sub() {
        let mut reg = Register::new(0xFF);
        let mut freg = Register::new(Flags::empty());

        sub(&mut reg, 0x0F, &mut freg);
        assert!(reg.read() == 0xF0);
        assert!(freg.read() == SubtractFlag);

        reg.write(0x11);
        sub(&mut reg, 0x11, &mut freg);
        assert!(reg.read() == 0);
        assert!(freg.read().contains(SubtractFlag));
        assert!(freg.read().contains(ZeroFlag));

        reg.write(0xA0);
        sub(&mut reg, 0xB0, &mut freg);
        assert!(reg.read() == 0xF0);
        assert!(freg.read() == SubtractFlag | CarryFlag);

        reg.write(0x80);
        sub(&mut reg, 0x0F, &mut freg);
        assert!(reg.read() == 0x71);
        assert!(freg.read() == SubtractFlag | HalfCarryFlag);

        reg.write(0xE5);
        sub(&mut reg, 0xAB, &mut freg);
        assert!(reg.read() == 0x3A);
        assert!(freg.read() == SubtractFlag | HalfCarryFlag);
    }

    #[test]
    fn test_decrement_value_at_address() {
        let mut mem = EmptyMemory::new(0xFFFF);
        let mut freg = Register::new(CarryFlag);
        mem.write_byte(0x1010, 1);

        decrement_value_at_address(&mut mem, 0x10, 0x10, &mut freg);

        assert!(mem.read_byte(0x1010) == 0);
        assert!(freg.read().is_all());

        mem.write_byte(0x01AB, 0x20);

        decrement_value_at_address(&mut mem, 0x01, 0xAB, &mut freg);

        assert!(mem.read_byte(0x01AB) == 0x1F);
        assert!(freg.read() == SubtractFlag);

        freg.write(ZeroFlag);
        mem.write_byte(0xABCD, 0xED);
        decrement_value_at_address(&mut mem, 0xAB, 0xCD, &mut freg);

        assert!(mem.read_byte(0xABCD) == 0xEC);
        assert!(freg.read() == SubtractFlag);

    }

    #[test]
    fn test_decrement_register_pair() {
        let mut reg1 = Register::new(0x70);
        let mut reg2 = Register::new(0x00);

        decrement_register_pair(&mut reg1, &mut reg2);

        assert!(reg1.read() == 0x6F);
        assert!(reg2.read() == 0xFF);
    }

    #[test]
    fn test_decrement_register() {
        let mut reg = Register::new(1);
        let mut freg = Register::new(Flags::empty());

        decrement_register(&mut reg, &mut freg);

        assert!(reg.read() == 0);
        assert!(freg.read() == ZeroFlag | SubtractFlag);

        reg.write(0x0);
        decrement_register(&mut reg, &mut freg);

        assert!(reg.read() == 255);
        assert!(freg.read() == ZeroFlag | SubtractFlag | CarryFlag | HalfCarryFlag);

        reg.write(0xF1);
        freg.write(Flags::empty());

        decrement_register(&mut reg, &mut freg);

        assert!(reg.read() == 0xF0);
        assert!(freg.read() == SubtractFlag);

        reg.write(0xF0);
        freg.write(Flags::empty());

        decrement_register(&mut reg, &mut freg);

        assert!(reg.read() == 0xEF);
        assert!(freg.read() == SubtractFlag);
    }
}
<commit_msg>Fix tests<commit_after>use memory::{EmptyMemory, Memory, pack_u16, high_byte, low_byte, low_nibble, high_nibble};
use extensions::Incrementor;
use cpu::{Register, Flags, CarryFlag, HalfCarryFlag, ZeroFlag, SubtractFlag};
use ops::u8_immediate;

fn half_carry_for_subtract(val1: u8, val2: u8, carry: u8) -> bool {
    return (low_nibble(val2) + carry) > low_nibble(val1);
}

pub fn sub_u8_immediate(mem: &Memory, pc: &mut Register<u16>, reg: &mut Register<u8>, freg: &mut Register<Flags>, with_carry: bool) {
    let val = u8_immediate(mem, pc);
    internal_sub(reg, val, freg, with_carry);
}

fn carry_for_subtract(val1: u8, val2: u8, carry: u8) -> bool {
    return (val2 as u16 + carry as u16) > val1 as u16;
}

pub fn internal_sub(reg: &mut Register<u8>, val: u8, freg: &mut Register<Flags>, with_carry: bool) {
    let reg_val = reg.read();
    let mut flags = SubtractFlag;
    let mut carry = 0;

    if freg.read().contains(CarryFlag) && with_carry {
        carry = 1;
    }

    if half_carry_for_subtract(reg_val, val, carry) {
        flags.insert(HalfCarryFlag);
    }

    if carry_for_subtract(reg_val, val, carry) {
        flags.insert(CarryFlag);
    }

    let result = reg_val.wrapping_sub(val).wrapping_sub(carry);

    if result == 0 {
        flags.insert(ZeroFlag);
    }
    debug!("internal sub result {:X}", result);
    reg.write(result);
    freg.write(flags);
}

pub fn sub(reg: &mut Register<u8>, val: u8, freg: &mut Register<Flags>) {
    internal_sub(reg, val, freg, false);
}

pub fn sbc(reg: &mut Register<u8>, val: u8, freg: &mut Register<Flags>) {
    internal_sub(reg, val, freg, true);
}

pub fn sub_value_at_address(mem: &Memory, reg: &mut Register<u8>, addr: u16, freg: &mut Register<Flags>) {
    let val = mem.read_byte(addr);
    internal_sub(reg, val, freg, false);
}

pub fn sbc_value_at_address(mem: &Memory, reg: &mut Register<u8>, addr: u16, freg: &mut Register<Flags>) {
    let val = mem.read_byte(addr);
    internal_sub(reg, val, freg, true);
}

pub fn compare(reg: &mut Register<u8>, val: u8, freg: &mut Register<Flags>) {
    let mut reg_val = Register::new(reg.read());
    internal_sub(&mut reg_val, val, freg, false);
}

pub fn compare_value_at_address(mem: &Memory, reg: &mut Register<u8>, addr: u16, freg: &mut Register<Flags>) {
    let val = mem.read_byte(addr);
    compare(reg, val, freg);
}

pub fn decrement_register(reg: &mut Register<u8>, freg: &mut Register<Flags>) {
    internal_sub(reg, 1, freg, false);
}

pub fn decrement_register_pair(reg1: &mut Register<u8>, reg2: &mut Register<u8>) {
    let val = pack_u16(reg1.read(), reg2.read());
    let ans = val - 1;

    reg1.write(high_byte(ans));
    reg2.write(low_byte(ans));
}

pub fn decrement_value_at_address(mem: &mut Memory, hb: u8, lb: u8, freg: &mut Register<Flags>) {
    let addr = pack_u16(hb, lb);
    let val = mem.read_byte(addr);
    let mut reg = Register::new(val);
    decrement_register(&mut reg, freg);
    mem.write_byte(addr, reg.read());
}

#[cfg(test)]
mod tests {
    use super::*;
    use memory::{EmptyMemory, Memory, pack_u16, high_byte, low_byte, low_nibble, high_nibble};
    use extensions::Incrementor;
    use cpu::{Register, Flags, CarryFlag, HalfCarryFlag, ZeroFlag, SubtractFlag};

    #[test]
    fn test_sub_u8_immediate() {
        let mut mem = EmptyMemory::new(0xFFFF);
        let mut pc = Register::new(0x2736);
        let mut reg = Register::new(0x35);
        let mut freg = Register::new(Flags::empty());
        mem.write_byte(pc.read(), 0x11);

        sub_u8_immediate(&mem, &mut pc, &mut reg, &mut freg, false);

        assert!(pc.read() == 0x2737);
        assert!(reg.read() == 0x24);
        assert!(freg.read() == SubtractFlag);

        mem.write_byte(pc.read(), 0x0C);
        sub_u8_immediate(&mem, &mut pc, &mut reg, &mut freg, false);

        assert!(pc.read() == 0x2738);
        assert!(reg.read() == 0x18);
        assert!(freg.read() == SubtractFlag | HalfCarryFlag);

        mem.write_byte(pc.read(), 0x0F);
        freg.write(CarryFlag);
        sub_u8_immediate(&mem, &mut pc, &mut reg, &mut freg, true);

        assert!(pc.read() == 0x2739);
        assert!(reg.read() == 0x08);
        assert!(freg.read() == SubtractFlag | HalfCarryFlag);
    }

    #[test]
    fn test_cp() {
        let mut reg = Register::new(0xAA);
        let mut val = 0xBB;
        let mut freg = Register::new(Flags::empty());

        compare(&mut reg, val, &mut freg);
        assert!(reg.read() == 0xAA);
        assert!(freg.read() == SubtractFlag | HalfCarryFlag | CarryFlag);

        reg.write(0xF0);
        val = 0xF0;
        compare(&mut reg, val, &mut freg);
        assert!(reg.read() == 0xF0);
        assert!(freg.read() == SubtractFlag | ZeroFlag);
    }

    #[test]
    fn test_compare_value_at_address() {
        let mut mem = EmptyMemory::new(0xFFFF);
        let mut reg = Register::new(0xBA);
        let mut freg = Register::new(Flags::empty());
        let addr = 0x1239;
        let mut val = 0xAB;
        mem.write_byte(addr, val);

        compare_value_at_address(&mem, &mut reg, addr, &mut freg);
        assert!(reg.read() == 0xBA);
        assert!(freg.read() == SubtractFlag | HalfCarryFlag);

        mem.write_byte(addr, 0xCD);
        reg.write(0xCD);
        compare_value_at_address(&mem, &mut reg, addr, &mut freg);
        assert!(reg.read() == 0xCD);
        assert!(freg.read() == SubtractFlag | ZeroFlag);

        mem.write_byte(addr, 0xCD);
        reg.write(0xCF);
        compare_value_at_address(&mem, &mut reg, addr, &mut freg);
        assert!(reg.read() == 0xCF);
        assert!(freg.read() == SubtractFlag);
    }

    #[test]
    fn test_sbc_value_at_address() {
        let mut mem = EmptyMemory::new(0xFFFF);
        let mut freg = Register::new(CarryFlag);
        let mut reg = Register::new(0xFF);
        let addr = 0x1234;
        mem.write_byte(addr, 0xCD);

        sbc_value_at_address(&mem, &mut reg, addr, &mut freg);
        assert!(reg.read() == 0x31);
        assert!(freg.read() == SubtractFlag);
    }

    #[test]
    fn test_sub_value_at_address() {
        let mut mem = EmptyMemory::new(0xFFFF);
        let mut freg = Register::new(CarryFlag);
        let mut reg = Register::new(0xFF);
        let addr = 0x1234;
        mem.write_byte(addr, 0xCD);

        sub_value_at_address(&mem, &mut reg, addr, &mut freg);
        assert!(reg.read() == 0x32);
        assert!(freg.read() == SubtractFlag);
    }

    #[test]
    fn test_sbc() {
        let mut reg = Register::new(0xFF);
        let mut freg = Register::new(Flags::empty());

        sbc(&mut reg, 0x0F, &mut freg);
        assert!(reg.read() == 0xF0);
        assert!(freg.read() == SubtractFlag);

        reg.write(0xFF);
        freg.write(CarryFlag);
        sbc(&mut reg, 0x0F, &mut freg);
        assert!(reg.read() == 0xEF);
        assert!(freg.read() == SubtractFlag | HalfCarryFlag);

        reg.write(0xFF);
        freg.write(CarryFlag);
        sbc(&mut reg, 0xFF, &mut freg);
        assert!(reg.read() == 0xFF);
        assert!(freg.read() == SubtractFlag | HalfCarryFlag | CarryFlag);

        reg.write(0xAB);
        freg.write(CarryFlag);
        sbc(&mut reg, 0x12, &mut freg);
        assert!(reg.read() == 0x98);
        assert!(freg.read() == SubtractFlag);
    }

    #[test]
    fn test_sub() {
        let mut reg = Register::new(0xFF);
        let mut freg = Register::new(Flags::empty());

        sub(&mut reg, 0x0F, &mut freg);
        assert!(reg.read() == 0xF0);
        assert!(freg.read() == SubtractFlag);

        reg.write(0x11);
        sub(&mut reg, 0x11, &mut freg);
        assert!(reg.read() == 0);
        assert!(freg.read().contains(SubtractFlag));
        assert!(freg.read().contains(ZeroFlag));

        reg.write(0xA0);
        sub(&mut reg, 0xB0, &mut freg);
        assert!(reg.read() == 0xF0);
        assert!(freg.read() == SubtractFlag | CarryFlag);

        reg.write(0x80);
        sub(&mut reg, 0x0F, &mut freg);
        assert!(reg.read() == 0x71);
        assert!(freg.read() == SubtractFlag | HalfCarryFlag);

        reg.write(0xE5);
        sub(&mut reg, 0xAB, &mut freg);
        assert!(reg.read() == 0x3A);
        assert!(freg.read() == SubtractFlag | HalfCarryFlag);
    }

    #[test]
    fn test_decrement_value_at_address() {
        let mut mem = EmptyMemory::new(0xFFFF);
        let mut freg = Register::new(CarryFlag);
        mem.write_byte(0x1010, 1);

        decrement_value_at_address(&mut mem, 0x10, 0x10, &mut freg);

        assert!(mem.read_byte(0x1010) == 0);
        assert!(freg.read() == SubtractFlag | ZeroFlag);

        mem.write_byte(0x01AB, 0x20);

        decrement_value_at_address(&mut mem, 0x01, 0xAB, &mut freg);

        assert!(mem.read_byte(0x01AB) == 0x1F);
        assert!(freg.read() == SubtractFlag | HalfCarryFlag);

        freg.write(ZeroFlag);
        mem.write_byte(0xABCD, 0xED);
        decrement_value_at_address(&mut mem, 0xAB, 0xCD, &mut freg);

        assert!(mem.read_byte(0xABCD) == 0xEC);
        assert!(freg.read() == SubtractFlag);

    }

    #[test]
    fn test_decrement_register_pair() {
        let mut reg1 = Register::new(0x70);
        let mut reg2 = Register::new(0x00);

        decrement_register_pair(&mut reg1, &mut reg2);

        assert!(reg1.read() == 0x6F);
        assert!(reg2.read() == 0xFF);
    }

    #[test]
    fn test_decrement_register() {
        let mut reg = Register::new(1);
        let mut freg = Register::new(Flags::empty());

        decrement_register(&mut reg, &mut freg);

        assert!(reg.read() == 0);
        assert!(freg.read() == ZeroFlag | SubtractFlag);

        reg.write(0x0);
        decrement_register(&mut reg, &mut freg);

        assert!(reg.read() == 0xFF);
        assert!(freg.read() == SubtractFlag | CarryFlag | HalfCarryFlag);

        reg.write(0xF1);
        freg.write(Flags::empty());

        decrement_register(&mut reg, &mut freg);

        assert!(reg.read() == 0xF0);
        assert!(freg.read() == SubtractFlag);

        reg.write(0xF0);
        freg.write(Flags::empty());

        decrement_register(&mut reg, &mut freg);

        assert!(reg.read() == 0xEF);
        assert!(freg.read() == SubtractFlag | HalfCarryFlag);
    }
}
<|endoftext|>


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

//! Synchronization of the local blockchain with the rest of the network. Used
//! either on a brand new node or when a node is late based on others' heads.
//! Always starts by downloading the header chain before asking either for full
//! blocks or a full UTXO set with related information.

/// How many block bodies to download in parallel
const MAX_BODY_DOWNLOADS: usize = 8;

use std::ops::{Deref, DerefMut};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::{Duration, Instant};

use core::core::hash::{Hash, Hashed};
use chain;
use p2p;
use types::Error;
use util::LOGGER;

#[derive(Debug)]
struct BlockDownload {
	hash: Hash,
	start_time: Instant,
	retries: u8,
}

/// Manages syncing the local chain with other peers. Needs both a head chain
/// and a full block chain to operate. First tries to advance the header
/// chain as much as possible, then downloads the full blocks by batches.
pub struct Syncer {
	chain: Arc<chain::Chain>,
	p2p: Arc<p2p::Server>,

	sync: Mutex<bool>,
	last_header_req: Mutex<Instant>,
	blocks_to_download: Mutex<Vec<Hash>>,
	blocks_downloading: Mutex<Vec<BlockDownload>>,
}

impl Syncer {
	pub fn new(chain_ref: Arc<chain::Chain>, p2p: Arc<p2p::Server>) -> Syncer {
		Syncer {
			chain: chain_ref,
			p2p: p2p,
			sync: Mutex::new(true),
			last_header_req: Mutex::new(Instant::now() - Duration::from_secs(2)),
			blocks_to_download: Mutex::new(vec![]),
			blocks_downloading: Mutex::new(vec![]),
		}
	}

	pub fn syncing(&self) -> bool {
		*self.sync.lock().unwrap()
	}

	/// Checks the local chain state, comparing it with our peers and triggers
	/// syncing if required.
	pub fn run(&self) -> Result<(), Error> {
		info!(LOGGER, "Sync: starting sync");

		// Loop for 10s waiting for some peers to potentially sync from.
		let start = Instant::now();
		loop {
			let pc = self.p2p.peer_count();
			if pc > 3 {
				break;
			}
			if Instant::now() - start > Duration::from_secs(10) {
				break;
			}
			thread::sleep(Duration::from_millis(200));
		}

		// Now check we actually have at least one peer to sync from.
		// If not then end the sync cleanly.
		if self.p2p.peer_count() == 0 {
			info!(LOGGER, "Sync: no peers to sync with, done.");

			let mut sync = self.sync.lock().unwrap();
			*sync = false;

			return Ok(())
		}

		// check if we have missing full blocks for which we already have a header
		self.init_download()?;

		// main syncing loop, requests more headers and bodies periodically as long
		// as a peer with higher difficulty exists and we're not fully caught up
		info!(LOGGER, "Sync: Starting loop.");
		loop {
			let tip = self.chain.get_header_head()?;

			// TODO do something better (like trying to get more) if we lose peers
			let peer = self.p2p.most_work_peer().expect("No peers available for sync.");
			let peer = peer.read().unwrap();
			debug!(
				LOGGER,
				"Sync: peer {} vs us {}",
				peer.info.total_difficulty,
				tip.total_difficulty
			);

			let more_headers = peer.info.total_difficulty > tip.total_difficulty;
			let more_bodies = {
				let blocks_to_download = self.blocks_to_download.lock().unwrap();
				let blocks_downloading = self.blocks_downloading.lock().unwrap();
				debug!(
					LOGGER,
					"Sync: blocks to download {}, block downloading {}",
					blocks_to_download.len(),
					blocks_downloading.len(),
				);
				blocks_to_download.len() > 0 || blocks_downloading.len() > 0
			};

			{
				let last_header_req = self.last_header_req.lock().unwrap().clone();
				if more_headers || (Instant::now() - Duration::from_secs(30) > last_header_req) {
					self.request_headers()?;
				}
			}
			if more_bodies {
				self.request_bodies();
			}
			if !more_headers && !more_bodies {
				// TODO check we haven't been lied to on the total work
				let mut sync = self.sync.lock().unwrap();
				*sync = false;
				break;
			}

			thread::sleep(Duration::from_secs(2));
		}
		info!(LOGGER, "Sync: done.");
		Ok(())
	}

	/// Checks the gap between the header chain and the full block chain and
	/// initializes the blocks_to_download structure with the missing full
	/// blocks
	fn init_download(&self) -> Result<(), Error> {
		// compare the header's head to the full one to see what we're missing
		let header_head = self.chain.get_header_head()?;
		let full_head = self.chain.head()?;
		let mut blocks_to_download = self.blocks_to_download.lock().unwrap();

		// go back the chain and insert for download all blocks we only have the
		// head for
		let mut prev_h = header_head.last_block_h;
		while prev_h != full_head.last_block_h {
			let header = self.chain.get_block_header(&prev_h)?;
			if header.height < full_head.height {
				break;
			}
			blocks_to_download.push(header.hash());
			prev_h = header.previous;
		}

		debug!(
			LOGGER,
			"Sync: Added {} full block hashes to download.",
			blocks_to_download.len()
		);
		Ok(())
	}

	/// Asks for the blocks we haven't downloaded yet and place them in the
	/// downloading structure.
	fn request_bodies(&self) {
		let mut blocks_to_download = self.blocks_to_download.lock().unwrap();
		let mut blocks_downloading = self.blocks_downloading.lock().unwrap();

		// retry blocks not downloading
		let now = Instant::now();
		for download in blocks_downloading.deref_mut() {
			let elapsed = (now - download.start_time).as_secs();
			if download.retries >= 12 {
				panic!("Failed to download required block {}", download.hash);
			}
			if download.retries < (elapsed / 5) as u8 {
				debug!(
					LOGGER,
					"Sync: Retry {} on block {}",
					download.retries,
					download.hash
				);
				self.request_block(download.hash);
				download.retries += 1;
			}
		}

		// consume hashes from blocks to download, place them in downloading and
		// request them from the network
		let mut count = 0;
		while blocks_to_download.len() > 0 && blocks_downloading.len() < MAX_BODY_DOWNLOADS {
			let h = blocks_to_download.pop().unwrap();
			self.request_block(h);
			count += 1;
			blocks_downloading.push(BlockDownload {
				hash: h,
				start_time: Instant::now(),
				retries: 0,
			});
		}
		debug!(
			LOGGER,
			"Sync: Requested {} full blocks to download, total left: {}. Current list: {:?}.",
			count,
			blocks_to_download.len(),
			blocks_downloading.deref(),
		);
	}

	/// We added a block, clean up the downloading structure
	pub fn block_received(&self, bh: Hash) {
		// just clean up the downloading list
		let mut bds = self.blocks_downloading.lock().unwrap();
		bds.iter()
			.position(|ref h| h.hash == bh)
			.map(|n| bds.remove(n));
	}

	/// Request some block headers from a peer to advance us
	fn request_headers(&self) -> Result<(), Error> {
		{
			let mut last_header_req = self.last_header_req.lock().unwrap();
			*last_header_req = Instant::now();
		}

		let tip = self.chain.get_header_head()?;
		let peer = self.p2p.most_work_peer();
		let locator = self.get_locator(&tip)?;
		if let Some(p) = peer {
			let p = p.read().unwrap();
			debug!(
				LOGGER,
				"Sync: Asking peer {} for more block headers, locator: {:?}",
				p.info.addr,
				locator,
			);
			if let Err(e) = p.send_header_request(locator) {
				debug!(LOGGER, "Sync: peer error, will retry");
			}
		} else {
			warn!(LOGGER, "Sync: Could not get most worked peer to request headers.");
		}
		Ok(())
	}

	/// We added a header, add it to the full block download list
	pub fn headers_received(&self, bhs: Vec<Hash>) {
		let mut blocks_to_download = self.blocks_to_download.lock().unwrap();
		for h in bhs {
			// enlist for full block download
			blocks_to_download.insert(0, h);
		}

		// we may still have more headers to retrieve but the main loop
		// will take care of this for us
	}

	/// Builds a vector of block hashes that should help the remote peer sending
	/// us the right block headers.
	fn get_locator(&self, tip: &chain::Tip) -> Result<Vec<Hash>, Error> {
		let heights = get_locator_heights(tip.height);

		debug!(LOGGER, "Sync: locator heights: {:?}", heights);

		let locator = heights
			.into_iter()
			.map(|h| {
				let header = self.chain.get_header_by_height(h).unwrap();
				header.hash()
			})
			.collect();

		debug!(LOGGER, "Sync: locator: {:?}", locator);

		Ok(locator)
	}

	/// Pick a random peer and ask for a block by hash
	fn request_block(&self, h: Hash) {
		if let Some(peer) = self.p2p.random_peer() {
			let peer = peer.read().unwrap();
			if let Err(e) = peer.send_block_request(h) {
				debug!(LOGGER, "Sync: Error requesting block: {:?}", e);
			}
		}
	}
}

// current height back to 0 decreasing in powers of 2
fn get_locator_heights(height: u64) -> Vec<u64> {
	let mut current = height.clone();
	let mut heights = vec![];
	while current > 0 {
		heights.push(current);
		let next = 2u64.pow(heights.len() as u32);
		current = if current > next {
			current - next
		} else {
			0
		}
	}
	heights.push(0);
	heights
}

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

	#[test]
	fn test_get_locator_heights() {
		assert_eq!(get_locator_heights(0), vec![0]);
		assert_eq!(get_locator_heights(1), vec![1, 0]);
		assert_eq!(get_locator_heights(2), vec![2, 0]);
		assert_eq!(get_locator_heights(3), vec![3, 1, 0]);
		assert_eq!(get_locator_heights(10), vec![10, 8, 4, 0]);
		assert_eq!(get_locator_heights(100), vec![100, 98, 94, 86, 70, 38, 0]);
		assert_eq!(
			get_locator_heights(1000),
			vec![1000, 998, 994, 986, 970, 938, 874, 746, 490, 0]
		);
		// check the locator is still a manageable length, even for large numbers of headers
		assert_eq!(
			get_locator_heights(10000),
			vec![10000, 9998, 9994, 9986, 9970, 9938, 9874, 9746, 9490, 8978, 7954, 5906, 1810, 0]
		);
	}
}
<commit_msg>Locator fix master (#381)<commit_after>// Copyright 2016 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Synchronization of the local blockchain with the rest of the network. Used
//! either on a brand new node or when a node is late based on others' heads.
//! Always starts by downloading the header chain before asking either for full
//! blocks or a full UTXO set with related information.

/// How many block bodies to download in parallel
const MAX_BODY_DOWNLOADS: usize = 8;

use std::ops::{Deref, DerefMut};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::{Duration, Instant};

use core::core::hash::{Hash, Hashed};
use chain;
use p2p;
use types::Error;
use util::LOGGER;

#[derive(Debug)]
struct BlockDownload {
	hash: Hash,
	start_time: Instant,
	retries: u8,
}

/// Manages syncing the local chain with other peers. Needs both a head chain
/// and a full block chain to operate. First tries to advance the header
/// chain as much as possible, then downloads the full blocks by batches.
pub struct Syncer {
	chain: Arc<chain::Chain>,
	p2p: Arc<p2p::Server>,

	sync: Mutex<bool>,
	last_header_req: Mutex<Instant>,
	blocks_to_download: Mutex<Vec<Hash>>,
	blocks_downloading: Mutex<Vec<BlockDownload>>,
}

impl Syncer {
	pub fn new(chain_ref: Arc<chain::Chain>, p2p: Arc<p2p::Server>) -> Syncer {
		Syncer {
			chain: chain_ref,
			p2p: p2p,
			sync: Mutex::new(true),
			last_header_req: Mutex::new(Instant::now() - Duration::from_secs(2)),
			blocks_to_download: Mutex::new(vec![]),
			blocks_downloading: Mutex::new(vec![]),
		}
	}

	pub fn syncing(&self) -> bool {
		*self.sync.lock().unwrap()
	}

	/// Checks the local chain state, comparing it with our peers and triggers
	/// syncing if required.
	pub fn run(&self) -> Result<(), Error> {
		info!(LOGGER, "Sync: starting sync");

		// Loop for 10s waiting for some peers to potentially sync from.
		let start = Instant::now();
		loop {
			let pc = self.p2p.peer_count();
			if pc > 3 {
				break;
			}
			if Instant::now() - start > Duration::from_secs(10) {
				break;
			}
			thread::sleep(Duration::from_millis(200));
		}

		// Now check we actually have at least one peer to sync from.
		// If not then end the sync cleanly.
		if self.p2p.peer_count() == 0 {
			info!(LOGGER, "Sync: no peers to sync with, done.");

			let mut sync = self.sync.lock().unwrap();
			*sync = false;

			return Ok(())
		}

		// check if we have missing full blocks for which we already have a header
		self.init_download()?;

		// main syncing loop, requests more headers and bodies periodically as long
		// as a peer with higher difficulty exists and we're not fully caught up
		info!(LOGGER, "Sync: Starting loop.");
		loop {
			let tip = self.chain.get_header_head()?;

			// TODO do something better (like trying to get more) if we lose peers
			let peer = self.p2p.most_work_peer().expect("No peers available for sync.");
			let peer = peer.read().unwrap();
			debug!(
				LOGGER,
				"Sync: peer {} vs us {}",
				peer.info.total_difficulty,
				tip.total_difficulty
			);

			let more_headers = peer.info.total_difficulty > tip.total_difficulty;
			let more_bodies = {
				let blocks_to_download = self.blocks_to_download.lock().unwrap();
				let blocks_downloading = self.blocks_downloading.lock().unwrap();
				debug!(
					LOGGER,
					"Sync: blocks to download {}, block downloading {}",
					blocks_to_download.len(),
					blocks_downloading.len(),
				);
				blocks_to_download.len() > 0 || blocks_downloading.len() > 0
			};

			{
				let last_header_req = self.last_header_req.lock().unwrap().clone();
				if more_headers || (Instant::now() - Duration::from_secs(30) > last_header_req) {
					self.request_headers()?;
				}
			}
			if more_bodies {
				self.request_bodies();
			}
			if !more_headers && !more_bodies {
				// TODO check we haven't been lied to on the total work
				let mut sync = self.sync.lock().unwrap();
				*sync = false;
				break;
			}

			thread::sleep(Duration::from_secs(2));
		}
		info!(LOGGER, "Sync: done.");
		Ok(())
	}

	/// Checks the gap between the header chain and the full block chain and
	/// initializes the blocks_to_download structure with the missing full
	/// blocks
	fn init_download(&self) -> Result<(), Error> {
		// compare the header's head to the full one to see what we're missing
		let header_head = self.chain.get_header_head()?;
		let full_head = self.chain.head()?;
		let mut blocks_to_download = self.blocks_to_download.lock().unwrap();

		// go back the chain and insert for download all blocks we only have the
		// head for
		let mut prev_h = header_head.last_block_h;
		while prev_h != full_head.last_block_h {
			let header = self.chain.get_block_header(&prev_h)?;
			if header.height < full_head.height {
				break;
			}
			blocks_to_download.push(header.hash());
			prev_h = header.previous;
		}

		debug!(
			LOGGER,
			"Sync: Added {} full block hashes to download.",
			blocks_to_download.len()
		);
		Ok(())
	}

	/// Asks for the blocks we haven't downloaded yet and place them in the
	/// downloading structure.
	fn request_bodies(&self) {
		let mut blocks_to_download = self.blocks_to_download.lock().unwrap();
		let mut blocks_downloading = self.blocks_downloading.lock().unwrap();

		// retry blocks not downloading
		let now = Instant::now();
		for download in blocks_downloading.deref_mut() {
			let elapsed = (now - download.start_time).as_secs();
			if download.retries >= 12 {
				panic!("Failed to download required block {}", download.hash);
			}
			if download.retries < (elapsed / 5) as u8 {
				debug!(
					LOGGER,
					"Sync: Retry {} on block {}",
					download.retries,
					download.hash
				);
				self.request_block(download.hash);
				download.retries += 1;
			}
		}

		// consume hashes from blocks to download, place them in downloading and
		// request them from the network
		let mut count = 0;
		while blocks_to_download.len() > 0 && blocks_downloading.len() < MAX_BODY_DOWNLOADS {
			let h = blocks_to_download.pop().unwrap();
			self.request_block(h);
			count += 1;
			blocks_downloading.push(BlockDownload {
				hash: h,
				start_time: Instant::now(),
				retries: 0,
			});
		}
		debug!(
			LOGGER,
			"Sync: Requested {} full blocks to download, total left: {}. Current list: {:?}.",
			count,
			blocks_to_download.len(),
			blocks_downloading.deref(),
		);
	}

	/// We added a block, clean up the downloading structure
	pub fn block_received(&self, bh: Hash) {
		// just clean up the downloading list
		let mut bds = self.blocks_downloading.lock().unwrap();
		bds.iter()
			.position(|ref h| h.hash == bh)
			.map(|n| bds.remove(n));
	}

	/// Request some block headers from a peer to advance us
	fn request_headers(&self) -> Result<(), Error> {
		{
			let mut last_header_req = self.last_header_req.lock().unwrap();
			*last_header_req = Instant::now();
		}

		let tip = self.chain.get_header_head()?;
		let peer = self.p2p.most_work_peer();
		let locator = self.get_locator(&tip)?;
		if let Some(p) = peer {
			let p = p.read().unwrap();
			debug!(
				LOGGER,
				"Sync: Asking peer {} for more block headers, locator: {:?}",
				p.info.addr,
				locator,
			);
			if let Err(e) = p.send_header_request(locator) {
				debug!(LOGGER, "Sync: peer error, will retry");
			}
		} else {
			warn!(LOGGER, "Sync: Could not get most worked peer to request headers.");
		}
		Ok(())
	}

	/// We added a header, add it to the full block download list
	pub fn headers_received(&self, bhs: Vec<Hash>) {
		let mut blocks_to_download = self.blocks_to_download.lock().unwrap();
		for h in bhs {
			// enlist for full block download
			blocks_to_download.insert(0, h);
		}

		// we may still have more headers to retrieve but the main loop
		// will take care of this for us
	}

	/// Builds a vector of block hashes that should help the remote peer sending
	/// us the right block headers.
	fn get_locator(&self, tip: &chain::Tip) -> Result<Vec<Hash>, Error> {
		let heights = get_locator_heights(tip.height);

		debug!(LOGGER, "Sync: locator heights: {:?}", heights);

		let locator = heights
			.into_iter()
			.map(|h| self.chain.get_header_by_height(h))
			.filter(|h| h.is_ok())
			.map(|h| h.unwrap().hash())
			.collect();
		debug!(LOGGER, "Sync: locator: {:?}", locator);
		Ok(locator)
	}

	/// Pick a random peer and ask for a block by hash
	fn request_block(&self, h: Hash) {
		if let Some(peer) = self.p2p.random_peer() {
			let peer = peer.read().unwrap();
			if let Err(e) = peer.send_block_request(h) {
				debug!(LOGGER, "Sync: Error requesting block: {:?}", e);
			}
		}
	}
}

// current height back to 0 decreasing in powers of 2
fn get_locator_heights(height: u64) -> Vec<u64> {
	let mut current = height.clone();
	let mut heights = vec![];
	while current > 0 {
		heights.push(current);
		let next = 2u64.pow(heights.len() as u32);
		current = if current > next {
			current - next
		} else {
			0
		}
	}
	heights.push(0);
	heights
}

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

	#[test]
	fn test_get_locator_heights() {
		assert_eq!(get_locator_heights(0), vec![0]);
		assert_eq!(get_locator_heights(1), vec![1, 0]);
		assert_eq!(get_locator_heights(2), vec![2, 0]);
		assert_eq!(get_locator_heights(3), vec![3, 1, 0]);
		assert_eq!(get_locator_heights(10), vec![10, 8, 4, 0]);
		assert_eq!(get_locator_heights(100), vec![100, 98, 94, 86, 70, 38, 0]);
		assert_eq!(
			get_locator_heights(1000),
			vec![1000, 998, 994, 986, 970, 938, 874, 746, 490, 0]
		);
		// check the locator is still a manageable length, even for large numbers of headers
		assert_eq!(
			get_locator_heights(10000),
			vec![10000, 9998, 9994, 9986, 9970, 9938, 9874, 9746, 9490, 8978, 7954, 5906, 1810, 0]
		);
	}
}
<|endoftext|>


<commit_before>use std::io::{Error, ErrorKind, Result};
use std::ffi::CString;
use message::Message;
use std::cmp::Ordering;
use std::ptr;
use std::boxed::Box;
use std::rc::Rc;
use std::ops::Deref;
use std::slice;
use std::time::Duration;

use ffi::*;

/// An LCM instance that handles publishing and subscribing,
/// as well as encoding and decoding messages.
pub struct Lcm {
    lcm: *mut lcm_t,
    subscriptions: Vec<Rc<LcmSubscription>>,
}

pub struct LcmSubscription {
    subscription: *mut lcm_subscription_t,
    handler: Box<FnMut(*const lcm_recv_buf_t)>,
}

impl Lcm {
    /// Creates a new `Lcm` instance.
    ///
    /// ```
    /// use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// ```
    pub fn new() -> Result<Lcm> {
        trace!("Creating LCM instance");
        let lcm = unsafe { lcm_create(ptr::null()) };
        match lcm.is_null() {
            true => Err(Error::new(ErrorKind::Other, "Failed to initialize LCM.")),
            false => {
                Ok(Lcm {
                    lcm: lcm,
                    subscriptions: Vec::new(),
                })
            }
        }
    }

    pub fn get_fileno(&self) -> ::std::os::raw::c_int {
        unsafe { lcm_get_fileno(self.lcm) }
    }

    /// Subscribes a callback to a particular topic.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// use std::sync::mpsc::channel;
    ///
    /// let mut lcm = Lcm::new().unwrap();
    /// let (tx, rx) = channel::<exlcm::Example>();
    /// lcm.subscribe("POSITION", |pos| { tx.send(pos).unwrap(); }
    /// ```
    pub fn subscribe<M, F>(&mut self, channel: &str, mut callback: F) -> Rc<LcmSubscription>
        where M: Message,
              F: FnMut(M) + 'static
    {
        trace!("Subscribing handler to channel {}", channel);

        let channel = CString::new(channel).unwrap();

        let handler = Box::new(move |rbuf: *const lcm_recv_buf_t| {
            trace!("Running handler");
            let mut buf = unsafe {
                let ref rbuf = *rbuf;
                let data = rbuf.data as *mut u8;
                let len = rbuf.data_size as usize;
                slice::from_raw_parts(data, len)
            };
            trace!("Decoding buffer: {:?}", buf);
            match M::decode_with_hash(&mut buf) {
                Ok(msg) => callback(msg),
                Err(_) => error!("Failed to decode buffer: {:?}", buf),
            }
        });

        let mut subscription = Rc::new(LcmSubscription {
            subscription: ptr::null_mut(),
            handler: handler,
        });

        let user_data = (subscription.deref() as *const _) as *mut _;

        let c_subscription = unsafe {
            lcm_subscribe(self.lcm,
                          channel.as_ptr(),
                          Some(Lcm::handler_callback::<M>),
                          user_data)
        };

        Rc::get_mut(&mut subscription).unwrap().subscription = c_subscription;
        self.subscriptions.push(subscription.clone());

        subscription
    }

    /// Unsubscribes a message handler.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// let handler = lcm.subscribe("POSITION", handler_function);
    /// // ...
    /// lcm.unsubscribe(handler);
    /// ```
    pub fn unsubscribe(&mut self, handler: LcmSubscription) -> Result<()> {
        trace!("Unsubscribing handler {:?}", handler.subscription);
        let result = unsafe { lcm_unsubscribe(self.lcm, handler.subscription) };
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM: Failed to unsubscribe")),
        }
    }

    /// Publishes a message on the specified channel.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    ///
    /// let mut my_data = exlcm::Example::new();
    /// my_data.timestamp = 0;
    /// my_data.position[0] = 1.0;
    /// my_data.position[1] = 2.0;
    /// my_data.position[2] = 3.0;
    ///
    /// lcm.publish("POSITION", &my_data).unwrap();
    /// ```
    pub fn publish<M>(&mut self, channel: &str, message: &M) -> Result<()>
        where M: Message
    {
        let channel = CString::new(channel).unwrap();
        let buffer = message.encode_with_hash()?;
        let result = unsafe {
            lcm_publish(self.lcm,
                        channel.as_ptr(),
                        buffer.as_ptr() as *mut _,
                        buffer.len() as _)
        };
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM Error")),
        }
    }

    /// Waits for and dispatches the next incoming message.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// lcm.subscribe("POSITION", handler_function);
    /// loop {
    ///     lcm.handle().unwrap();
    /// }
    /// ```
    pub fn handle(&mut self) -> Result<()> {
        let result = unsafe { lcm_handle(self.lcm) };
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM Error")),
        }
    }

    /// Waits for and dispatches the next incoming message, up to a time limit.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// lcm.subscribe("POSITION", handler_function);
    /// loop {
    ///     lcm.handle_timeout(1000).unwrap();
    /// }
    /// ```
    pub fn handle_timeout(&mut self, timeout: &Duration) -> Result<()> {
        let millis = (timeout.as_secs() * 1000) as i32 + (timeout.subsec_nanos() / 1_000_000) as i32;
        let result = unsafe { lcm_handle_timeout(self.lcm, millis) };
        match result.cmp(&0) {
            Ordering::Less => Err(Error::new(ErrorKind::Other, "LCM Error")),
            Ordering::Equal => Err(Error::new(ErrorKind::Other, "LCM Timeout")),
            Ordering::Greater => Ok(()),
        }
    }

    /// Adjusts the maximum number of received messages that can be queued up for a subscription.
    /// The default is `30`.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// let handler = lcm.subscribe("POSITION", handler_function);
    /// lcm.subscription_set_queue_capacity(handler, 30);
    /// ```
    pub fn subscription_set_queue_capacity(handler: Rc<LcmSubscription>, num_messages: usize) {
        let handler = handler.subscription;
        let num_messages = num_messages as _;
        unsafe { lcm_subscription_set_queue_capacity(handler, num_messages) };
    }

    extern "C" fn handler_callback<M>(rbuf: *const lcm_recv_buf_t,
                                      _: *const ::std::os::raw::c_char,
                                      user_data: *mut ::std::os::raw::c_void)
        where M: Message
    {
        trace!("Received data");
        let sub = user_data as *mut LcmSubscription;
        let sub = unsafe { &mut *sub };
        (sub.handler)(rbuf);
    }
}

impl Drop for Lcm {
    fn drop(&mut self) {
        trace!("Destroying Lcm instance");
        unsafe { lcm_destroy(self.lcm) };
    }
}
<commit_msg>Doctests now pass<commit_after>use std::io::{Error, ErrorKind, Result};
use std::ffi::CString;
use message::Message;
use std::cmp::Ordering;
use std::ptr;
use std::boxed::Box;
use std::rc::Rc;
use std::ops::Deref;
use std::slice;
use std::time::Duration;

use ffi::*;

/// An LCM instance that handles publishing and subscribing,
/// as well as encoding and decoding messages.
pub struct Lcm {
    lcm: *mut lcm_t,
    subscriptions: Vec<Rc<LcmSubscription>>,
}

pub struct LcmSubscription {
    subscription: *mut lcm_subscription_t,
    handler: Box<FnMut(*const lcm_recv_buf_t)>,
}

impl Lcm {
    /// Creates a new `Lcm` instance.
    ///
    /// ```
    /// use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// ```
    pub fn new() -> Result<Lcm> {
        trace!("Creating LCM instance");
        let lcm = unsafe { lcm_create(ptr::null()) };
        match lcm.is_null() {
            true => Err(Error::new(ErrorKind::Other, "Failed to initialize LCM.")),
            false => {
                Ok(Lcm {
                    lcm: lcm,
                    subscriptions: Vec::new(),
                })
            }
        }
    }

    pub fn get_fileno(&self) -> ::std::os::raw::c_int {
        unsafe { lcm_get_fileno(self.lcm) }
    }

    /// Subscribes a callback to a particular topic.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// lcm.subscribe("GREETINGS", |name: String| println!("Hello, {}!", name) );
    /// ```
    pub fn subscribe<M, F>(&mut self, channel: &str, mut callback: F) -> Rc<LcmSubscription>
        where M: Message,
              F: FnMut(M) + 'static
    {
        trace!("Subscribing handler to channel {}", channel);

        let channel = CString::new(channel).unwrap();

        let handler = Box::new(move |rbuf: *const lcm_recv_buf_t| {
            trace!("Running handler");
            let mut buf = unsafe {
                let ref rbuf = *rbuf;
                let data = rbuf.data as *mut u8;
                let len = rbuf.data_size as usize;
                slice::from_raw_parts(data, len)
            };
            trace!("Decoding buffer: {:?}", buf);
            match M::decode_with_hash(&mut buf) {
                Ok(msg) => callback(msg),
                Err(_) => error!("Failed to decode buffer: {:?}", buf),
            }
        });

        let mut subscription = Rc::new(LcmSubscription {
            subscription: ptr::null_mut(),
            handler: handler,
        });

        let user_data = (subscription.deref() as *const _) as *mut _;

        let c_subscription = unsafe {
            lcm_subscribe(self.lcm,
                          channel.as_ptr(),
                          Some(Lcm::handler_callback::<M>),
                          user_data)
        };

        Rc::get_mut(&mut subscription).unwrap().subscription = c_subscription;
        self.subscriptions.push(subscription.clone());

        subscription
    }

    /// Unsubscribes a message handler.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// # let handler_function = |name: String| println!("Hello, {}!", name);
    /// # let mut lcm = Lcm::new().unwrap();
    /// let handler = lcm.subscribe("GREETINGS", handler_function);
    /// // ...
    /// lcm.unsubscribe(handler);
    /// ```
    pub fn unsubscribe(&mut self, handler: Rc<LcmSubscription>) -> Result<()> {
        trace!("Unsubscribing handler {:?}", handler.subscription);
        let result = unsafe { lcm_unsubscribe(self.lcm, handler.subscription) };
        // TODO: Remove subscription from Lcm object
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM: Failed to unsubscribe")),
        }
    }

    /// Publishes a message on the specified channel.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// lcm.publish("GREETINGS", &"Charles".to_string()).unwrap();
    /// ```
    pub fn publish<M>(&mut self, channel: &str, message: &M) -> Result<()>
        where M: Message
    {
        let channel = CString::new(channel).unwrap();
        let buffer = message.encode_with_hash()?;
        let result = unsafe {
            lcm_publish(self.lcm,
                        channel.as_ptr(),
                        buffer.as_ptr() as *mut _,
                        buffer.len() as _)
        };
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM Error")),
        }
    }

    /// Waits for and dispatches the next incoming message.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// # let handler_function = |name: String| println!("Hello, {}!", name);
    /// let mut lcm = Lcm::new().unwrap();
    /// lcm.subscribe("POSITION", handler_function);
    /// loop {
    /// # break;
    ///     lcm.handle().unwrap();
    /// }
    /// ```
    pub fn handle(&mut self) -> Result<()> {
        let result = unsafe { lcm_handle(self.lcm) };
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM Error")),
        }
    }

    /// Waits for and dispatches the next incoming message, up to a time limit.
    ///
    /// ```
    /// # use std::time::Duration;
    /// # use lcm::Lcm;
    /// # let handler_function = |name: String| println!("Hello, {}!", name);
    /// let mut lcm = Lcm::new().unwrap();
    /// lcm.subscribe("POSITION", handler_function);
    /// let wait_dur = Duration::from_millis(100);
    /// loop {
    /// # break;
    ///     lcm.handle_timeout(&wait_dur).unwrap();
    /// }
    /// ```
    pub fn handle_timeout(&mut self, timeout: &Duration) -> Result<()> {
        let millis = (timeout.as_secs() * 1000) as i32 + (timeout.subsec_nanos() / 1_000_000) as i32;
        let result = unsafe { lcm_handle_timeout(self.lcm, millis) };
        match result.cmp(&0) {
            Ordering::Less => Err(Error::new(ErrorKind::Other, "LCM Error")),
            Ordering::Equal => Err(Error::new(ErrorKind::Other, "LCM Timeout")),
            Ordering::Greater => Ok(()),
        }
    }

    /// Adjusts the maximum number of received messages that can be queued up for a subscription.
    /// The default is `30`.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// # let handler_function = |name: String| println!("Hello, {}!", name);
    /// # let mut lcm = Lcm::new().unwrap();
    /// let handler = lcm.subscribe("POSITION", handler_function);
    /// lcm.subscription_set_queue_capacity(handler, 30);
    /// ```
    pub fn subscription_set_queue_capacity(&self, handler: Rc<LcmSubscription>, num_messages: usize) {
        let handler = handler.subscription;
        let num_messages = num_messages as _;
        unsafe { lcm_subscription_set_queue_capacity(handler, num_messages) };
    }

    extern "C" fn handler_callback<M>(rbuf: *const lcm_recv_buf_t,
                                      _: *const ::std::os::raw::c_char,
                                      user_data: *mut ::std::os::raw::c_void)
        where M: Message
    {
        trace!("Received data");
        let sub = user_data as *mut LcmSubscription;
        let sub = unsafe { &mut *sub };
        (sub.handler)(rbuf);
    }
}

impl Drop for Lcm {
    fn drop(&mut self) {
        trace!("Destroying Lcm instance");
        unsafe { lcm_destroy(self.lcm) };
    }
}
<|endoftext|>


<commit_before>//! # The Rust Core Library
//!
//! The Rust Core Library is the dependency-free[^free] foundation of [The
//! Rust Standard Library](../std/index.html). It is the portable glue
//! between the language and its libraries, defining the intrinsic and
//! primitive building blocks of all Rust code. It links to no
//! upstream libraries, no system libraries, and no libc.
//!
//! [^free]: Strictly speaking, there are some symbols which are needed but
//!          they aren't always necessary.
//!
//! The core library is *minimal*: it isn't even aware of heap allocation,
//! nor does it provide concurrency or I/O. These things require
//! platform integration, and this library is platform-agnostic.
//!
//! # How to use the core library
//!
//! Please note that all of these details are currently not considered stable.
//!
// FIXME: Fill me in with more detail when the interface settles
//! This library is built on the assumption of a few existing symbols:
//!
//! * `memcpy`, `memcmp`, `memset` - These are core memory routines which are
//!   often generated by LLVM. Additionally, this library can make explicit
//!   calls to these functions. Their signatures are the same as found in C.
//!   These functions are often provided by the system libc, but can also be
//!   provided by the [compiler-builtins crate](https://crates.io/crates/compiler_builtins).
//!
//! * `rust_begin_panic` - This function takes four arguments, a
//!   `fmt::Arguments`, a `&'static str`, and two `u32`'s. These four arguments
//!   dictate the panic message, the file at which panic was invoked, and the
//!   line and column inside the file. It is up to consumers of this core
//!   library to define this panic function; it is only required to never
//!   return. This requires a `lang` attribute named `panic_impl`.
//!
//! * `rust_eh_personality` - is used by the failure mechanisms of the
//!    compiler. This is often mapped to GCC's personality function, but crates
//!    which do not trigger a panic can be assured that this function is never
//!    called. The `lang` attribute is called `eh_personality`.

// Since libcore defines many fundamental lang items, all tests live in a
// separate crate, libcoretest, to avoid bizarre issues.
//
// Here we explicitly #[cfg]-out this whole crate when testing. If we don't do
// this, both the generated test artifact and the linked libtest (which
// transitively includes libcore) will both define the same set of lang items,
// and this will cause the E0152 "found duplicate lang item" error. See
// discussion in #50466 for details.
//
// This cfg won't affect doc tests.
#![cfg(not(test))]
// To run libcore tests without x.py without ending up with two copies of libcore, Miri needs to be
// able to "empty" this crate. See <https://github.com/rust-lang/miri-test-libstd/issues/4>.
// rustc itself never sets the feature, so this line has no affect there.
#![cfg(any(not(feature = "miri-test-libstd"), test, doctest))]
#![stable(feature = "core", since = "1.6.0")]
#![doc(
    html_playground_url = "https://play.rust-lang.org/",
    issue_tracker_base_url = "https://github.com/rust-lang/rust/issues/",
    test(no_crate_inject, attr(deny(warnings))),
    test(attr(allow(dead_code, deprecated, unused_variables, unused_mut)))
)]
#![no_core]

// Lints:
#![deny(rust_2021_incompatible_or_patterns)]
#![deny(unsafe_op_in_unsafe_fn)]
#![warn(deprecated_in_future)]
#![warn(missing_debug_implementations)]
#![warn(missing_docs)]
#![allow(explicit_outlives_requirements)]

// Library features:
#![feature(const_align_of_val)]
#![feature(const_alloc_layout)]
#![feature(const_arguments_as_str)]
#![feature(const_assert_type)]
#![feature(const_caller_location)]
#![feature(const_cell_into_inner)]
#![feature(const_discriminant)]
#![feature(const_float_bits_conv)]
#![feature(const_float_classify)]
#![feature(const_heap)]
#![feature(const_inherent_unchecked_arith)]
#![feature(const_int_unchecked_arith)]
#![feature(const_intrinsic_copy)]
#![feature(const_intrinsic_forget)]
#![feature(const_likely)]
#![feature(const_maybe_uninit_as_ptr)]
#![feature(const_maybe_uninit_assume_init)]
#![feature(const_option)]
#![feature(const_pin)]
#![feature(const_ptr_offset)]
#![feature(const_ptr_offset_from)]
#![feature(const_ptr_read)]
#![feature(const_ptr_write)]
#![feature(const_raw_ptr_comparison)]
#![feature(const_size_of_val)]
#![feature(const_slice_from_raw_parts)]
#![feature(const_slice_ptr_len)]
#![feature(const_swap)]
#![feature(const_type_id)]
#![feature(const_type_name)]
#![feature(const_unreachable_unchecked)]
#![feature(duration_consts_2)]
#![feature(ptr_metadata)]
#![feature(slice_ptr_get)]
#![feature(std_internals)]
#![feature(variant_count)]

// Language features:
#![feature(abi_unadjusted)]
#![feature(allow_internal_unstable)]
#![feature(arbitrary_self_types)]
#![feature(asm)]
#![feature(associated_type_bounds)]
#![feature(auto_traits)]
#![feature(cfg_target_has_atomic)]
#![feature(const_fn_floating_point_arithmetic)]
#![feature(const_fn_fn_ptr_basics)]
#![feature(const_fn_trait_bound)]
#![cfg_attr(bootstrap, feature(const_fn_transmute))]
#![cfg_attr(bootstrap, feature(const_fn_union))]
#![feature(const_impl_trait)]
#![feature(const_mut_refs)]
#![feature(const_panic)]
#![feature(const_precise_live_drops)]
#![feature(const_raw_ptr_deref)]
#![feature(const_refs_to_cell)]
#![feature(custom_inner_attributes)]
#![feature(decl_macro)]
#![feature(doc_cfg)]
#![feature(doc_notable_trait)]
#![feature(exhaustive_patterns)]
#![feature(extern_types)]
#![feature(fundamental)]
#![feature(intra_doc_pointers)]
#![feature(intrinsics)]
#![feature(lang_items)]
#![feature(link_llvm_intrinsics)]
#![feature(llvm_asm)]
#![feature(min_specialization)]
#![feature(negative_impls)]
#![feature(never_type)]
#![feature(nll)]
#![feature(no_core)]
#![feature(no_coverage)] // rust-lang/rust#84605
#![feature(no_niche)] // rust-lang/rust#68303
#![feature(platform_intrinsics)]
#![feature(prelude_import)]
#![feature(repr_simd)]
#![feature(rustc_allow_const_fn_unstable)]
#![feature(rustc_attrs)]
#![feature(simd_ffi)]
#![feature(staged_api)]
#![feature(stmt_expr_attributes)]
#![feature(trait_alias)]
#![feature(transparent_unions)]
#![feature(try_blocks)]
#![feature(unboxed_closures)]
#![feature(unsized_fn_params)]

// Target features:
#![feature(aarch64_target_feature)]
#![feature(adx_target_feature)]
#![feature(arm_target_feature)]
#![feature(avx512_target_feature)]
#![feature(cmpxchg16b_target_feature)]
#![feature(f16c_target_feature)]
#![feature(hexagon_target_feature)]
#![feature(mips_target_feature)]
#![feature(powerpc_target_feature)]
#![feature(rtm_target_feature)]
#![feature(sse4a_target_feature)]
#![feature(tbm_target_feature)]
#![feature(wasm_target_feature)]

// allow using `core::` in intra-doc links
#[allow(unused_extern_crates)]
extern crate self as core;

#[prelude_import]
#[allow(unused)]
use prelude::v1::*;

#[cfg(not(test))] // See #65860
#[macro_use]
mod macros;

// We don't export this through #[macro_export] for now, to avoid breakage.
// See https://github.com/rust-lang/rust/issues/82913
#[cfg(not(test))]
#[unstable(feature = "assert_matches", issue = "82775")]
/// Unstable module containing the unstable `assert_matches` macro.
pub mod assert_matches {
    #[unstable(feature = "assert_matches", issue = "82775")]
    pub use crate::macros::{assert_matches, debug_assert_matches};
}

#[macro_use]
mod internal_macros;

#[path = "num/shells/int_macros.rs"]
#[macro_use]
mod int_macros;

#[path = "num/shells/i128.rs"]
pub mod i128;
#[path = "num/shells/i16.rs"]
pub mod i16;
#[path = "num/shells/i32.rs"]
pub mod i32;
#[path = "num/shells/i64.rs"]
pub mod i64;
#[path = "num/shells/i8.rs"]
pub mod i8;
#[path = "num/shells/isize.rs"]
pub mod isize;

#[path = "num/shells/u128.rs"]
pub mod u128;
#[path = "num/shells/u16.rs"]
pub mod u16;
#[path = "num/shells/u32.rs"]
pub mod u32;
#[path = "num/shells/u64.rs"]
pub mod u64;
#[path = "num/shells/u8.rs"]
pub mod u8;
#[path = "num/shells/usize.rs"]
pub mod usize;

#[path = "num/f32.rs"]
pub mod f32;
#[path = "num/f64.rs"]
pub mod f64;

#[macro_use]
pub mod num;

/* The libcore prelude, not as all-encompassing as the libstd prelude */

pub mod prelude;

/* Core modules for ownership management */

pub mod hint;
pub mod intrinsics;
pub mod mem;
pub mod ptr;

/* Core language traits */

pub mod borrow;
pub mod clone;
pub mod cmp;
pub mod convert;
pub mod default;
pub mod marker;
pub mod ops;

/* Core types and methods on primitives */

pub mod any;
pub mod array;
pub mod ascii;
pub mod cell;
pub mod char;
pub mod ffi;
pub mod iter;
#[unstable(feature = "once_cell", issue = "74465")]
pub mod lazy;
pub mod option;
pub mod panic;
pub mod panicking;
pub mod pin;
pub mod result;
#[unstable(feature = "async_stream", issue = "79024")]
pub mod stream;
pub mod sync;

pub mod fmt;
pub mod hash;
pub mod slice;
pub mod str;
pub mod time;

pub mod unicode;

/* Async */
pub mod future;
pub mod task;

/* Heap memory allocator trait */
#[allow(missing_docs)]
pub mod alloc;

// note: does not need to be public
mod bool;
mod tuple;
mod unit;

#[stable(feature = "core_primitive", since = "1.43.0")]
pub mod primitive;

// Pull in the `core_arch` crate directly into libcore. The contents of
// `core_arch` are in a different repository: rust-lang/stdarch.
//
// `core_arch` depends on libcore, but the contents of this module are
// set up in such a way that directly pulling it here works such that the
// crate uses the this crate as its libcore.
#[path = "../../stdarch/crates/core_arch/src/mod.rs"]
#[allow(
    missing_docs,
    missing_debug_implementations,
    dead_code,
    unused_imports,
    unsafe_op_in_unsafe_fn
)]
#[allow(rustdoc::bare_urls)]
// FIXME: This annotation should be moved into rust-lang/stdarch after clashing_extern_declarations is
// merged. It currently cannot because bootstrap fails as the lint hasn't been defined yet.
#[allow(clashing_extern_declarations)]
#[unstable(feature = "stdsimd", issue = "48556")]
mod core_arch;

#[doc = include_str!("../../stdarch/crates/core_arch/src/core_arch_docs.md")]
#[stable(feature = "simd_arch", since = "1.27.0")]
pub mod arch {
    #[stable(feature = "simd_arch", since = "1.27.0")]
    pub use crate::core_arch::arch::*;

    /// Inline assembly.
    ///
    /// Read the [unstable book] for the usage.
    ///
    /// [unstable book]: ../../unstable-book/library-features/asm.html
    #[unstable(
        feature = "asm",
        issue = "72016",
        reason = "inline assembly is not stable enough for use and is subject to change"
    )]
    #[rustc_builtin_macro]
    pub macro asm("assembly template", $(operands,)* $(options($(option),*))?) {
        /* compiler built-in */
    }

    /// Module-level inline assembly.
    #[unstable(
        feature = "global_asm",
        issue = "35119",
        reason = "`global_asm!` is not stable enough for use and is subject to change"
    )]
    #[rustc_builtin_macro]
    pub macro global_asm("assembly template", $(operands,)* $(options($(option),*))?) {
        /* compiler built-in */
    }
}
<commit_msg>Remove unused langauge #![feature]s from core.<commit_after>//! # The Rust Core Library
//!
//! The Rust Core Library is the dependency-free[^free] foundation of [The
//! Rust Standard Library](../std/index.html). It is the portable glue
//! between the language and its libraries, defining the intrinsic and
//! primitive building blocks of all Rust code. It links to no
//! upstream libraries, no system libraries, and no libc.
//!
//! [^free]: Strictly speaking, there are some symbols which are needed but
//!          they aren't always necessary.
//!
//! The core library is *minimal*: it isn't even aware of heap allocation,
//! nor does it provide concurrency or I/O. These things require
//! platform integration, and this library is platform-agnostic.
//!
//! # How to use the core library
//!
//! Please note that all of these details are currently not considered stable.
//!
// FIXME: Fill me in with more detail when the interface settles
//! This library is built on the assumption of a few existing symbols:
//!
//! * `memcpy`, `memcmp`, `memset` - These are core memory routines which are
//!   often generated by LLVM. Additionally, this library can make explicit
//!   calls to these functions. Their signatures are the same as found in C.
//!   These functions are often provided by the system libc, but can also be
//!   provided by the [compiler-builtins crate](https://crates.io/crates/compiler_builtins).
//!
//! * `rust_begin_panic` - This function takes four arguments, a
//!   `fmt::Arguments`, a `&'static str`, and two `u32`'s. These four arguments
//!   dictate the panic message, the file at which panic was invoked, and the
//!   line and column inside the file. It is up to consumers of this core
//!   library to define this panic function; it is only required to never
//!   return. This requires a `lang` attribute named `panic_impl`.
//!
//! * `rust_eh_personality` - is used by the failure mechanisms of the
//!    compiler. This is often mapped to GCC's personality function, but crates
//!    which do not trigger a panic can be assured that this function is never
//!    called. The `lang` attribute is called `eh_personality`.

// Since libcore defines many fundamental lang items, all tests live in a
// separate crate, libcoretest, to avoid bizarre issues.
//
// Here we explicitly #[cfg]-out this whole crate when testing. If we don't do
// this, both the generated test artifact and the linked libtest (which
// transitively includes libcore) will both define the same set of lang items,
// and this will cause the E0152 "found duplicate lang item" error. See
// discussion in #50466 for details.
//
// This cfg won't affect doc tests.
#![cfg(not(test))]
// To run libcore tests without x.py without ending up with two copies of libcore, Miri needs to be
// able to "empty" this crate. See <https://github.com/rust-lang/miri-test-libstd/issues/4>.
// rustc itself never sets the feature, so this line has no affect there.
#![cfg(any(not(feature = "miri-test-libstd"), test, doctest))]
#![stable(feature = "core", since = "1.6.0")]
#![doc(
    html_playground_url = "https://play.rust-lang.org/",
    issue_tracker_base_url = "https://github.com/rust-lang/rust/issues/",
    test(no_crate_inject, attr(deny(warnings))),
    test(attr(allow(dead_code, deprecated, unused_variables, unused_mut)))
)]
#![no_core]

// Lints:
#![deny(rust_2021_incompatible_or_patterns)]
#![deny(unsafe_op_in_unsafe_fn)]
#![warn(deprecated_in_future)]
#![warn(missing_debug_implementations)]
#![warn(missing_docs)]
#![allow(explicit_outlives_requirements)]

// Library features:
#![feature(const_align_of_val)]
#![feature(const_alloc_layout)]
#![feature(const_arguments_as_str)]
#![feature(const_assert_type)]
#![feature(const_caller_location)]
#![feature(const_cell_into_inner)]
#![feature(const_discriminant)]
#![feature(const_float_bits_conv)]
#![feature(const_float_classify)]
#![feature(const_heap)]
#![feature(const_inherent_unchecked_arith)]
#![feature(const_int_unchecked_arith)]
#![feature(const_intrinsic_copy)]
#![feature(const_intrinsic_forget)]
#![feature(const_likely)]
#![feature(const_maybe_uninit_as_ptr)]
#![feature(const_maybe_uninit_assume_init)]
#![feature(const_option)]
#![feature(const_pin)]
#![feature(const_ptr_offset)]
#![feature(const_ptr_offset_from)]
#![feature(const_ptr_read)]
#![feature(const_ptr_write)]
#![feature(const_raw_ptr_comparison)]
#![feature(const_size_of_val)]
#![feature(const_slice_from_raw_parts)]
#![feature(const_slice_ptr_len)]
#![feature(const_swap)]
#![feature(const_type_id)]
#![feature(const_type_name)]
#![feature(const_unreachable_unchecked)]
#![feature(duration_consts_2)]
#![feature(ptr_metadata)]
#![feature(slice_ptr_get)]
#![feature(std_internals)]
#![feature(variant_count)]

// Language features:
#![feature(abi_unadjusted)]
#![feature(allow_internal_unstable)]
#![feature(asm)]
#![feature(associated_type_bounds)]
#![feature(auto_traits)]
#![feature(cfg_target_has_atomic)]
#![feature(const_fn_floating_point_arithmetic)]
#![feature(const_fn_fn_ptr_basics)]
#![feature(const_fn_trait_bound)]
#![cfg_attr(bootstrap, feature(const_fn_transmute))]
#![cfg_attr(bootstrap, feature(const_fn_union))]
#![feature(const_impl_trait)]
#![feature(const_mut_refs)]
#![feature(const_panic)]
#![feature(const_precise_live_drops)]
#![feature(const_raw_ptr_deref)]
#![feature(const_refs_to_cell)]
#![feature(decl_macro)]
#![feature(doc_cfg)]
#![feature(doc_notable_trait)]
#![feature(exhaustive_patterns)]
#![feature(extern_types)]
#![feature(fundamental)]
#![feature(intra_doc_pointers)]
#![feature(intrinsics)]
#![feature(lang_items)]
#![feature(link_llvm_intrinsics)]
#![feature(llvm_asm)]
#![feature(min_specialization)]
#![feature(negative_impls)]
#![feature(never_type)]
#![feature(no_core)]
#![feature(no_coverage)] // rust-lang/rust#84605
#![feature(no_niche)] // rust-lang/rust#68303
#![feature(platform_intrinsics)]
#![feature(prelude_import)]
#![feature(repr_simd)]
#![feature(rustc_allow_const_fn_unstable)]
#![feature(rustc_attrs)]
#![feature(simd_ffi)]
#![feature(staged_api)]
#![feature(stmt_expr_attributes)]
#![feature(trait_alias)]
#![feature(transparent_unions)]
#![feature(try_blocks)]
#![feature(unboxed_closures)]
#![feature(unsized_fn_params)]

// Target features:
#![feature(aarch64_target_feature)]
#![feature(adx_target_feature)]
#![feature(arm_target_feature)]
#![feature(avx512_target_feature)]
#![feature(cmpxchg16b_target_feature)]
#![feature(f16c_target_feature)]
#![feature(hexagon_target_feature)]
#![feature(mips_target_feature)]
#![feature(powerpc_target_feature)]
#![feature(rtm_target_feature)]
#![feature(sse4a_target_feature)]
#![feature(tbm_target_feature)]
#![feature(wasm_target_feature)]

// allow using `core::` in intra-doc links
#[allow(unused_extern_crates)]
extern crate self as core;

#[prelude_import]
#[allow(unused)]
use prelude::v1::*;

#[cfg(not(test))] // See #65860
#[macro_use]
mod macros;

// We don't export this through #[macro_export] for now, to avoid breakage.
// See https://github.com/rust-lang/rust/issues/82913
#[cfg(not(test))]
#[unstable(feature = "assert_matches", issue = "82775")]
/// Unstable module containing the unstable `assert_matches` macro.
pub mod assert_matches {
    #[unstable(feature = "assert_matches", issue = "82775")]
    pub use crate::macros::{assert_matches, debug_assert_matches};
}

#[macro_use]
mod internal_macros;

#[path = "num/shells/int_macros.rs"]
#[macro_use]
mod int_macros;

#[path = "num/shells/i128.rs"]
pub mod i128;
#[path = "num/shells/i16.rs"]
pub mod i16;
#[path = "num/shells/i32.rs"]
pub mod i32;
#[path = "num/shells/i64.rs"]
pub mod i64;
#[path = "num/shells/i8.rs"]
pub mod i8;
#[path = "num/shells/isize.rs"]
pub mod isize;

#[path = "num/shells/u128.rs"]
pub mod u128;
#[path = "num/shells/u16.rs"]
pub mod u16;
#[path = "num/shells/u32.rs"]
pub mod u32;
#[path = "num/shells/u64.rs"]
pub mod u64;
#[path = "num/shells/u8.rs"]
pub mod u8;
#[path = "num/shells/usize.rs"]
pub mod usize;

#[path = "num/f32.rs"]
pub mod f32;
#[path = "num/f64.rs"]
pub mod f64;

#[macro_use]
pub mod num;

/* The libcore prelude, not as all-encompassing as the libstd prelude */

pub mod prelude;

/* Core modules for ownership management */

pub mod hint;
pub mod intrinsics;
pub mod mem;
pub mod ptr;

/* Core language traits */

pub mod borrow;
pub mod clone;
pub mod cmp;
pub mod convert;
pub mod default;
pub mod marker;
pub mod ops;

/* Core types and methods on primitives */

pub mod any;
pub mod array;
pub mod ascii;
pub mod cell;
pub mod char;
pub mod ffi;
pub mod iter;
#[unstable(feature = "once_cell", issue = "74465")]
pub mod lazy;
pub mod option;
pub mod panic;
pub mod panicking;
pub mod pin;
pub mod result;
#[unstable(feature = "async_stream", issue = "79024")]
pub mod stream;
pub mod sync;

pub mod fmt;
pub mod hash;
pub mod slice;
pub mod str;
pub mod time;

pub mod unicode;

/* Async */
pub mod future;
pub mod task;

/* Heap memory allocator trait */
#[allow(missing_docs)]
pub mod alloc;

// note: does not need to be public
mod bool;
mod tuple;
mod unit;

#[stable(feature = "core_primitive", since = "1.43.0")]
pub mod primitive;

// Pull in the `core_arch` crate directly into libcore. The contents of
// `core_arch` are in a different repository: rust-lang/stdarch.
//
// `core_arch` depends on libcore, but the contents of this module are
// set up in such a way that directly pulling it here works such that the
// crate uses the this crate as its libcore.
#[path = "../../stdarch/crates/core_arch/src/mod.rs"]
#[allow(
    missing_docs,
    missing_debug_implementations,
    dead_code,
    unused_imports,
    unsafe_op_in_unsafe_fn
)]
#[allow(rustdoc::bare_urls)]
// FIXME: This annotation should be moved into rust-lang/stdarch after clashing_extern_declarations is
// merged. It currently cannot because bootstrap fails as the lint hasn't been defined yet.
#[allow(clashing_extern_declarations)]
#[unstable(feature = "stdsimd", issue = "48556")]
mod core_arch;

#[doc = include_str!("../../stdarch/crates/core_arch/src/core_arch_docs.md")]
#[stable(feature = "simd_arch", since = "1.27.0")]
pub mod arch {
    #[stable(feature = "simd_arch", since = "1.27.0")]
    pub use crate::core_arch::arch::*;

    /// Inline assembly.
    ///
    /// Read the [unstable book] for the usage.
    ///
    /// [unstable book]: ../../unstable-book/library-features/asm.html
    #[unstable(
        feature = "asm",
        issue = "72016",
        reason = "inline assembly is not stable enough for use and is subject to change"
    )]
    #[rustc_builtin_macro]
    pub macro asm("assembly template", $(operands,)* $(options($(option),*))?) {
        /* compiler built-in */
    }

    /// Module-level inline assembly.
    #[unstable(
        feature = "global_asm",
        issue = "35119",
        reason = "`global_asm!` is not stable enough for use and is subject to change"
    )]
    #[rustc_builtin_macro]
    pub macro global_asm("assembly template", $(operands,)* $(options($(option),*))?) {
        /* compiler built-in */
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
use std::net::{UdpSocket, TcpStream};
use std::str;

// Content of the request.
const SEARCH_REQUEST: &'static str =
"M-SEARCH * HTTP/1.1\r
Host:239.255.255.250:1900\r
ST:urn:schemas-upnp-org:device:InternetGatewayDevice:1\r
Man:\"ssdp:discover\"\r
MX:3\r\n\r\n";

// Error type this module emits.
#[derive(Debug)]
pub enum SearchError {
    IoError(io::Error),
    InvalidResponse,
}

impl From<io::Error> for SearchError {
    fn from(err: io::Error) -> SearchError {
        SearchError::IoError(err)
    }
}

// Try to find the gateway on the local network.
pub fn search_gateway() -> Result<SocketAddrV4, SearchError> {
    let local_ip = try!(get_local_ip());
    let local_addr = (local_ip, 1900);
    let socket = try!(UdpSocket::bind(local_addr));

    // send the request on the broadcast address
    try!(socket.send_to(SEARCH_REQUEST.as_bytes(), "239.255.255.250:1900"));

    let mut buf = [0u8; 1024];
    let (read, _) = try!(socket.recv_from(&mut buf));
    let text = str::from_utf8(&buf[..read]).unwrap();
    match parse_result(text) {
        None => Err(SearchError::InvalidResponse),
        Some(socketaddr) => Ok(socketaddr),
    }
}

// This is a hacky way of getting this computer's address
// on the local network. There doesn't seem to be another
// way of doing this at the moment.
fn get_local_ip() -> io::Result<Ipv4Addr> {
    let sock = try!(TcpStream::connect("google.ca:80"));
    let sockaddr = try!(sock.local_addr());
    match sockaddr {
        SocketAddr::V4(ref v4) => Ok(v4.ip().clone()),
        _ => panic!("unsupported"),
    }
}

// Parse the result.
fn parse_result(text: &str) -> Option<SocketAddrV4> {
    let re = regex!(r"LOCATION:\s*http://(\d+\.\d+\.\d+\.\d+):(\d+)/.+");
    for line in text.lines() {
        match re.captures(line) {
            None => continue,
            Some(cap) => {
                // these shouldn't fail if the regex matched.
                let addr = cap.at(1).unwrap();
                let port = cap.at(2).unwrap();
                return Some(SocketAddrV4::new(
                    addr.parse::<Ipv4Addr>().unwrap(),
                    port.parse::<u16>().unwrap()
                ));
            },
        }
    }
    None
}
<commit_msg>add function to let the user to choose network interface to be used<commit_after>use std::io;
use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
use std::net::{UdpSocket, TcpStream};
use std::str;

// Content of the request.
const SEARCH_REQUEST: &'static str =
"M-SEARCH * HTTP/1.1\r
Host:239.255.255.250:1900\r
ST:urn:schemas-upnp-org:device:InternetGatewayDevice:1\r
Man:\"ssdp:discover\"\r
MX:3\r\n\r\n";

// Error type this module emits.
#[derive(Debug)]
pub enum SearchError {
    IoError(io::Error),
    InvalidResponse,
}

impl From<io::Error> for SearchError {
    fn from(err: io::Error) -> SearchError {
        SearchError::IoError(err)
    }
}

pub fn search_gateway_from(local_addr: SocketAddrV4) -> Result<SocketAddrV4, SearchError> {
    let socket = try!(UdpSocket::bind(local_addr));

    // send the request on the broadcast address
    try!(socket.send_to(SEARCH_REQUEST.as_bytes(), "239.255.255.250:1900"));

    let mut buf = [0u8; 1024];
    let (read, _) = try!(socket.recv_from(&mut buf));
    let text = str::from_utf8(&buf[..read]).unwrap();
    match parse_result(text) {
        None => Err(SearchError::InvalidResponse),
        Some(socketaddr) => Ok(socketaddr),
    }
}

// Try to find the gateway on the local network.
pub fn search_gateway() -> Result<SocketAddrV4, SearchError> {
    let local_ip = try!(get_local_ip());
    search_gateway_from(SocketAddrV4::new(local_ip, 1900))
}

// This is a hacky way of getting this computer's address
// on the local network. There doesn't seem to be another
// way of doing this at the moment.
fn get_local_ip() -> io::Result<Ipv4Addr> {
    let sock = try!(TcpStream::connect("google.ca:80"));
    let sockaddr = try!(sock.local_addr());
    match sockaddr {
        SocketAddr::V4(ref v4) => Ok(v4.ip().clone()),
        _ => panic!("unsupported"),
    }
}

// Parse the result.
fn parse_result(text: &str) -> Option<SocketAddrV4> {
    let re = regex!(r"LOCATION:\s*http://(\d+\.\d+\.\d+\.\d+):(\d+)/.+");
    for line in text.lines() {
        match re.captures(line) {
            None => continue,
            Some(cap) => {
                // these shouldn't fail if the regex matched.
                let addr = cap.at(1).unwrap();
                let port = cap.at(2).unwrap();
                return Some(SocketAddrV4::new(
                    addr.parse::<Ipv4Addr>().unwrap(),
                    port.parse::<u16>().unwrap()
                ));
            },
        }
    }
    None
}
<|endoftext|>