<commit_before>/*
 * CMOS IO handling
 */

use vm;

use std::rc::Rc;
use std::cell::RefCell;
use time;

const CMOS_SELECT_PORT: u16     = 0x70;
const CMOS_DATA_PORT: u16       = 0x71;
const CMOS_TOTAL_REGS: u8       = 128;  // Total number of byte registers we emulate
const CMOS_DEFAULT_SELECTOR: u8 = 0xD;  // Default selected register

/*
 * Handled registers:
 * 0x00      Seconds
 * 0x02      Minutes
 * 0x04      Hours
 * 0x06      Weekday
 * 0x07      Day of Month
 * 0x08      Month
 * 0x09      Year
 * 0x32      Century (maybe)
 * 0x0A      Status Register A
 * 0x0B      Status Register B
 * 0x0D      Status Register D
 *
 * Current limitations:
 * - No BCD format, only binary mode
 * - No 12 hour support, only 24
 * - No interrupt generation
 * - RTC updates on every register access which can lead to unstable time readings in guests
 */
struct CMOS
{
    selector: u8,
    sta: u8,
    stb: u8,
    nmi_bit: bool,          // TODO: this bit should be owned by vm/vcpu
    host_time: time::Tm,    // Real host time during last update
    time: time::Tm,         // Time we are emulating
}

impl CMOS
{
    fn new() -> CMOS
    {
        CMOS {
            selector: CMOS_DEFAULT_SELECTOR,
            sta: 0b00100110,
            stb: 0b00000110,
            nmi_bit: false,
            host_time: time::empty_tm(),
            time: time::empty_tm(),
        }
    }

    fn nmi_mask(&self) -> u8
    {
        if self.nmi_bit {
            return 0x80;
        } else {
            return 0;
        }
    }

    fn read_selector(&mut self) -> u8
    {
        self.selector | self.nmi_mask()
    }

    fn write_selector(&mut self, val: u8)
    {
        self.selector = val & 0x7F;
        self.nmi_bit = (val & 0x80) != 0;
    }

    // Returns current selector value and resets it to default
    fn reset_selector(&mut self) -> u8
    {
        let val = self.selector;
        self.selector = CMOS_DEFAULT_SELECTOR;
        return val;
    }

    fn to_rtc_format(&self, val: i32) -> u8
    {
        if (self.stb & 0x04) == 0 {
            // BCD format needed
            assert!(val < 100);
            let lo = (val % 10) as u8;
            let hi = (val / 10) as u8;
            return lo | (hi << 4);
        } else {
            return val as u8;
        }
    }

    fn from_rtc_format(&self, val: u8) -> i32
    {
        if (self.stb & 0x04) == 0 {
            // BCD format needed
            let lo = (val & 0xF) as i32;
            let hi = (val >> 4) as i32;
            return lo + hi;
        } else {
            return val as i32;
        }
    }

    fn read_reg(&mut self) -> u8
    {
        // Adjust emulated time by computing elapsed duration since last time
        // Then add this duration to time we emulate
        let now = time::now();
        let delta = now - self.host_time; // Ok if negative

        self.time = self.time + delta;
        self.host_time = now;

        return match self.reset_selector() {
            // RTC
            0x00 => self.to_rtc_format(self.time.tm_sec),
            0x02 => self.to_rtc_format(self.time.tm_min),
            0x04 => self.to_rtc_format(self.time.tm_hour),
            0x06 => self.to_rtc_format(self.time.tm_wday + 1), // CMOS wday starts from 1
            0x07 => self.to_rtc_format(self.time.tm_mday),
            0x08 => self.to_rtc_format(self.time.tm_mon),
            0x09 => self.to_rtc_format(self.time.tm_year % 100),
            0x32 => self.to_rtc_format(self.time.tm_year / 100),

            // Status
            0x0A => self.sta,
            0x0B => self.stb,

            // Unsupported
            _ => 0,
        } as u8;
    }

    fn write_reg(&mut self, val: u8)
    {
        match self.reset_selector() {
            // RTC
            0x00 => self.time.tm_sec    = self.from_rtc_format(val),
            0x02 => self.time.tm_min    = self.from_rtc_format(val),
            0x04 => self.time.tm_hour   = self.from_rtc_format(val),
            0x06 => self.time.tm_wday   = self.from_rtc_format(val) - 1, // CMOS wday starts from 1
            0x07 => self.time.tm_mday   = self.from_rtc_format(val),
            0x08 => self.time.tm_mon    = self.from_rtc_format(val),
            0x09 => self.time.tm_year   = self.from_rtc_format(val) + self.time.tm_year / 100 * 100,
            0x32 => self.time.tm_year   = self.from_rtc_format(val) * 100 + self.time.tm_year % 100,

            // Status
            0x0A => self.sta = val,
            0x0B => self.stb = val,

            // Unsupported
            _ => (),
        };
    }
}

#[cfg(test)]
mod cmos_test {

    use super::CMOS;
    use time;

    fn read_reg(cmos: &mut CMOS, reg: u8) -> u8
    {
        let sel = (cmos.read_selector() & 0x80) | reg;
        cmos.write_selector(sel);
        return cmos.read_reg();
    }

    fn write_reg(cmos: &mut CMOS, reg: u8, val: u8)
    {
        let sel = (cmos.read_selector() & 0x80) | reg;
        cmos.write_selector(sel);
        cmos.write_reg(val);
    }

    fn read_rtc_reg(cmos: &mut CMOS, is_bcd: bool, reg: u8) -> u8
    {
        let mut val = read_reg(cmos, reg);

        if is_bcd {
            let lo = val & 0xF;
            let hi = val >> 4;
            val = hi * 10 + lo;
        }

        return val;

    }

    fn set_bcd(cmos: &mut CMOS, is_bcd: bool)
    {
        let stb = read_reg(cmos, 0x0b);
        if is_bcd {
            write_reg(cmos, 0x0b, stb | 0x04u8);
        } else {
            write_reg(cmos, 0x0b, stb & !0x04u8);
        }
    }

    // Test initial state
    #[test] fn default()
    {
        let cmos = CMOS::new();

        assert!(cmos.selector == super::CMOS_DEFAULT_SELECTOR);
        assert!(cmos.nmi_bit == false);
        assert!(cmos.host_time == time::empty_tm());
        assert!(cmos.time == time::empty_tm());
    }

    // Check that NMI bit is propogated to selector value
    #[test] fn nmi_bit()
    {
        let mut cmos = CMOS::new();

        let mut sel = cmos.read_selector();
        assert!(sel & 0x80 == 0);

        cmos.write_selector(sel | 0x80);
        sel = cmos.read_selector();
        assert!(sel & 0x80 != 0);

        cmos.write_selector(sel & 0x7F);
        sel = cmos.read_selector();
        assert!(sel & 0x80 == 0);
    }

    fn rtc_to_tm(cmos: &mut CMOS) -> time::Tm
    {
        let is_bcd = (read_reg(cmos, 0x0b) & 0x04) == 0;

        let mut tm: time::Tm = time::empty_tm();
        tm.tm_sec = read_rtc_reg(cmos, is_bcd, 0x00) as i32;
        tm.tm_min = read_rtc_reg(cmos, is_bcd, 0x02) as i32;
        tm.tm_hour = read_rtc_reg(cmos, is_bcd, 0x04) as i32;
        tm.tm_wday = read_rtc_reg(cmos, is_bcd, 0x06) as i32;
        tm.tm_mday = read_rtc_reg(cmos, is_bcd, 0x07) as i32;
        tm.tm_mon = read_rtc_reg(cmos, is_bcd, 0x08) as i32;
        tm.tm_year = (read_rtc_reg(cmos, is_bcd, 0x32) as i32) * 100 + (read_rtc_reg(cmos, is_bcd, 0x09) as i32);

        return tm;
    }

    fn rtc_test_common(cmos: &mut CMOS)
    {
        // Time reported by rtc should be within sane distance from now
        let now = time::now();
        let t = rtc_to_tm(cmos);
        assert!((t - now).num_seconds() <= 1);

        // Check that eventually rtc will report a monotonic increase in time
        loop {
            let tm1 = rtc_to_tm(cmos);
            let tm2 = rtc_to_tm(cmos);
            if tm2 > tm1 {
                break;
            }
        }
    }

    #[test] fn rtc()
    {
        let mut cmos = CMOS::new();

        set_bcd(&mut cmos, true);
        rtc_test_common(&mut cmos);

        set_bcd(&mut cmos, false);
        rtc_test_common(&mut cmos);
    }
}

// TODO: MOAR TESTS

///////////////////////////////////////////////////////////////////////////////

struct CMOSDev
{
    cmos: RefCell<CMOS>,
}

impl vm::io_handler for CMOSDev
{
    fn io_read(&self, port: u16, size: u8) -> vm::IoOperandType
    {
        let mut cmos = self.cmos.borrow_mut();

        assert!(size == 1);
        assert!(cmos.selector < CMOS_TOTAL_REGS);

        match port {
            CMOS_SELECT_PORT => {
                return vm::IoOperandType::byte(cmos.read_selector());
            },

            CMOS_DATA_PORT => {
                return vm::IoOperandType::byte(cmos.read_reg());
            }

            _ => {
                panic!();
            }
        }
    }

    fn io_write(&self, port: u16, data: vm::IoOperandType)
    {
        let mut cmos = self.cmos.borrow_mut();
        let val: u8 = data.unwrap_byte();

        assert!(cmos.selector < CMOS_TOTAL_REGS);

        match port {
            CMOS_SELECT_PORT => {
                cmos.write_selector(val);
            }

            CMOS_DATA_PORT => {
                cmos.write_reg(val);
            }

            _ => {
                panic!();
            }
        }
    }
}

pub fn init(vm: &mut vm::vm)
{
	let dev = Rc::new(CMOSDev {
        cmos: RefCell::new(CMOS::new()),
    });

    vm::register_io_region(vm, dev.clone(), CMOS_SELECT_PORT, 1);
    vm::register_io_region(vm, dev.clone(), CMOS_DATA_PORT, 1);
}

<commit_msg>cmos: check unsupported status bits<commit_after>/*
 * CMOS IO handling
 */

use vm;

use std::rc::Rc;
use std::cell::RefCell;
use time;

const CMOS_SELECT_PORT: u16     = 0x70;
const CMOS_DATA_PORT: u16       = 0x71;
const CMOS_TOTAL_REGS: u8       = 128;  // Total number of byte registers we emulate
const CMOS_DEFAULT_SELECTOR: u8 = 0xD;  // Default selected register
const CMOS_STA_DEFAULT: u8      = 0b00100110;
const CMOS_STA_SUPPORTED: u8    = 0b00000000;
const CMOS_STB_DEFAULT: u8      = 0b00000110;
const CMOS_STB_SUPPORTED: u8    = 0b00000100;

/*
 * Handled registers:
 * 0x00      Seconds
 * 0x02      Minutes
 * 0x04      Hours
 * 0x06      Weekday
 * 0x07      Day of Month
 * 0x08      Month
 * 0x09      Year
 * 0x32      Century (maybe)
 * 0x0A      Status Register A
 * 0x0B      Status Register B
 * 0x0D      Status Register D
 *
 * Current limitations:
 * - No 12 hour support, only 24
 * - No interrupt generation
 * - RTC updates on every register access which can lead to unstable time readings in guests
 */
struct CMOS
{
    selector: u8,
    sta: u8,
    stb: u8,
    nmi_bit: bool,          // TODO: this bit should be owned by vm/vcpu
    host_time: time::Tm,    // Real host time during last update
    time: time::Tm,         // Time we are emulating
}

impl CMOS
{
    fn new() -> CMOS
    {
        CMOS {
            selector: CMOS_DEFAULT_SELECTOR,
            sta: CMOS_STA_DEFAULT,
            stb: CMOS_STB_DEFAULT,
            nmi_bit: false,
            host_time: time::empty_tm(),
            time: time::empty_tm(),
        }
    }

    fn nmi_mask(&self) -> u8
    {
        if self.nmi_bit {
            return 0x80;
        } else {
            return 0;
        }
    }

    fn read_selector(&mut self) -> u8
    {
        self.selector | self.nmi_mask()
    }

    fn write_selector(&mut self, val: u8)
    {
        self.selector = val & 0x7F;
        self.nmi_bit = (val & 0x80) != 0;
    }

    // Returns current selector value and resets it to default
    fn reset_selector(&mut self) -> u8
    {
        let val = self.selector;
        self.selector = CMOS_DEFAULT_SELECTOR;
        return val;
    }

    fn to_rtc_format(&self, val: i32) -> u8
    {
        if (self.stb & 0x04) == 0 {
            // BCD format needed
            assert!(val < 100);
            let lo = (val % 10) as u8;
            let hi = (val / 10) as u8;
            return lo | (hi << 4);
        } else {
            return val as u8;
        }
    }

    fn from_rtc_format(&self, val: u8) -> i32
    {
        if (self.stb & 0x04) == 0 {
            // BCD format needed
            let lo = (val & 0xF) as i32;
            let hi = (val >> 4) as i32;
            return lo + hi;
        } else {
            return val as i32;
        }
    }

    fn read_reg(&mut self) -> u8
    {
        // Adjust emulated time by computing elapsed duration since last time
        // Then add this duration to time we emulate
        let now = time::now();
        let delta = now - self.host_time; // Ok if negative

        self.time = self.time + delta;
        self.host_time = now;

        return match self.reset_selector() {
            // RTC
            0x00 => self.to_rtc_format(self.time.tm_sec),
            0x02 => self.to_rtc_format(self.time.tm_min),
            0x04 => self.to_rtc_format(self.time.tm_hour),
            0x06 => self.to_rtc_format(self.time.tm_wday + 1), // CMOS wday starts from 1
            0x07 => self.to_rtc_format(self.time.tm_mday),
            0x08 => self.to_rtc_format(self.time.tm_mon),
            0x09 => self.to_rtc_format(self.time.tm_year % 100),
            0x32 => self.to_rtc_format(self.time.tm_year / 100),

            // Status
            0x0A => self.sta,
            0x0B => self.stb,

            // Unsupported
            _ => 0,
        } as u8;
    }

    fn write_reg(&mut self, val: u8)
    {
        match self.reset_selector() {
            // RTC
            0x00 => self.time.tm_sec    = self.from_rtc_format(val),
            0x02 => self.time.tm_min    = self.from_rtc_format(val),
            0x04 => self.time.tm_hour   = self.from_rtc_format(val),
            0x06 => self.time.tm_wday   = self.from_rtc_format(val) - 1, // CMOS wday starts from 1
            0x07 => self.time.tm_mday   = self.from_rtc_format(val),
            0x08 => self.time.tm_mon    = self.from_rtc_format(val),
            0x09 => self.time.tm_year   = self.from_rtc_format(val) + self.time.tm_year / 100 * 100,
            0x32 => self.time.tm_year   = self.from_rtc_format(val) * 100 + self.time.tm_year % 100,

            // Status
            0x0A => {
                let diff = self.sta ^ val;
                if (diff & !CMOS_STA_SUPPORTED) != 0 {
                    panic!("CMOS: setting unsupported STA bits");
                }
                self.sta = val;
            },

            0x0B => {
                let diff = self.stb ^ val;
                if (diff & !CMOS_STB_SUPPORTED) != 0 {
                    panic!("CMOS: setting unsupported STB bits");
                }
                self.stb = val;
            },

            // Unsupported
            _ => (),
        };
    }
}

#[cfg(test)]
mod cmos_test {

    use super::CMOS;
    use time;

    fn read_reg(cmos: &mut CMOS, reg: u8) -> u8
    {
        let sel = (cmos.read_selector() & 0x80) | reg;
        cmos.write_selector(sel);
        return cmos.read_reg();
    }

    fn write_reg(cmos: &mut CMOS, reg: u8, val: u8)
    {
        let sel = (cmos.read_selector() & 0x80) | reg;
        cmos.write_selector(sel);
        cmos.write_reg(val);
    }

    fn read_rtc_reg(cmos: &mut CMOS, is_bcd: bool, reg: u8) -> u8
    {
        let mut val = read_reg(cmos, reg);

        if is_bcd {
            let lo = val & 0xF;
            let hi = val >> 4;
            val = hi * 10 + lo;
        }

        return val;

    }

    fn set_bcd(cmos: &mut CMOS, is_bcd: bool)
    {
        let stb = read_reg(cmos, 0x0b);
        if is_bcd {
            write_reg(cmos, 0x0b, stb | 0x04);
        } else {
            write_reg(cmos, 0x0b, stb & !0x04);
        }
    }

    // Test initial state
    #[test] fn default()
    {
        let cmos = CMOS::new();

        assert!(cmos.selector == super::CMOS_DEFAULT_SELECTOR);
        assert!(cmos.nmi_bit == false);
        assert!(cmos.host_time == time::empty_tm());
        assert!(cmos.time == time::empty_tm());
    }

    // Check that NMI bit is propogated to selector value
    #[test] fn nmi_bit()
    {
        let mut cmos = CMOS::new();

        let mut sel = cmos.read_selector();
        assert!(sel & 0x80 == 0);

        cmos.write_selector(sel | 0x80);
        sel = cmos.read_selector();
        assert!(sel & 0x80 != 0);

        cmos.write_selector(sel & 0x7F);
        sel = cmos.read_selector();
        assert!(sel & 0x80 == 0);
    }

    fn rtc_to_tm(cmos: &mut CMOS) -> time::Tm
    {
        let is_bcd = (read_reg(cmos, 0x0b) & 0x04) == 0;

        let mut tm: time::Tm = time::empty_tm();
        tm.tm_sec = read_rtc_reg(cmos, is_bcd, 0x00) as i32;
        tm.tm_min = read_rtc_reg(cmos, is_bcd, 0x02) as i32;
        tm.tm_hour = read_rtc_reg(cmos, is_bcd, 0x04) as i32;
        tm.tm_wday = read_rtc_reg(cmos, is_bcd, 0x06) as i32;
        tm.tm_mday = read_rtc_reg(cmos, is_bcd, 0x07) as i32;
        tm.tm_mon = read_rtc_reg(cmos, is_bcd, 0x08) as i32;
        tm.tm_year = (read_rtc_reg(cmos, is_bcd, 0x32) as i32) * 100 + (read_rtc_reg(cmos, is_bcd, 0x09) as i32);

        return tm;
    }

    fn rtc_test_common(cmos: &mut CMOS)
    {
        // Time reported by rtc should be within sane distance from now
        let now = time::now();
        let t = rtc_to_tm(cmos);
        assert!((t - now).num_seconds() <= 1);
    }

    #[test] fn rtc()
    {
        let mut cmos = CMOS::new();

        set_bcd(&mut cmos, true);
        rtc_test_common(&mut cmos);

        set_bcd(&mut cmos, false);
        rtc_test_common(&mut cmos);

        // Check that eventually rtc will report a monotonic increase in time
        loop {
            let tm1 = rtc_to_tm(&mut cmos);
            let tm2 = rtc_to_tm(&mut cmos);
            if tm2 > tm1 {
                break;
            }
        }
    }
}

// TODO: MOAR TESTS

///////////////////////////////////////////////////////////////////////////////

struct CMOSDev
{
    cmos: RefCell<CMOS>,
}

impl vm::io_handler for CMOSDev
{
    fn io_read(&self, port: u16, size: u8) -> vm::IoOperandType
    {
        let mut cmos = self.cmos.borrow_mut();

        assert!(size == 1);
        assert!(cmos.selector < CMOS_TOTAL_REGS);

        match port {
            CMOS_SELECT_PORT => {
                return vm::IoOperandType::byte(cmos.read_selector());
            },

            CMOS_DATA_PORT => {
                return vm::IoOperandType::byte(cmos.read_reg());
            }

            _ => {
                panic!();
            }
        }
    }

    fn io_write(&self, port: u16, data: vm::IoOperandType)
    {
        let mut cmos = self.cmos.borrow_mut();
        let val: u8 = data.unwrap_byte();

        assert!(cmos.selector < CMOS_TOTAL_REGS);

        match port {
            CMOS_SELECT_PORT => {
                cmos.write_selector(val);
            }

            CMOS_DATA_PORT => {
                cmos.write_reg(val);
            }

            _ => {
                panic!();
            }
        }
    }
}

pub fn init(vm: &mut vm::vm)
{
	let dev = Rc::new(CMOSDev {
        cmos: RefCell::new(CMOS::new()),
    });

    vm::register_io_region(vm, dev.clone(), CMOS_SELECT_PORT, 1);
    vm::register_io_region(vm, dev.clone(), CMOS_DATA_PORT, 1);
}

<|endoftext|>


<commit_before>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://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.

#[cfg(feature = "wasm")]
pub mod js_fetcher;
pub mod mock_fetcher;

use crate::http::{HttpRequest, HttpResponse};
use crate::utils::{MaybeSend, MaybeSync};
use anyhow::{anyhow, Result};
use async_trait::async_trait;

/// An interface for fetching resources from network.
#[cfg_attr(feature = "wasm", async_trait(?Send))]
#[cfg_attr(not(feature = "wasm"), async_trait)]
pub trait Fetcher: MaybeSend + MaybeSync {
    async fn fetch(&self, request: HttpRequest) -> Result<HttpResponse>;
}

/// Uses `Get` method and returns response body.
/// - Why this function is not put inside [`Fetcher`] trait?
///   If we declare `Fetcher::get` function with a default implementation,
///   we have to also add a constraint `where Self: Sized` to `Fetcher::get`,
///   because of https://github.com/rust-lang/rust/issues/51443 and in particular
///   https://docs.rs/async-trait/0.1.57/async_trait/#dyn-traits.
///   However, having such constraint `Self: Sized` prevent using `Fetcher::get` method on a
///   `dyn Fetcher` variable, because `dyn Fetcher` is not `Sized`.
pub async fn get(fetcher: &dyn Fetcher, url: impl ToString) -> Result<Vec<u8>> {
    let request = HttpRequest {
        body: vec![],
        headers: vec![],
        method: crate::http::Method::Get,
        url: url.to_string(),
    };
    let response = fetcher.fetch(request).await?;
    Ok(response.body)
}

pub const NULL_FETCHER: NullFetcher = NullFetcher {};

pub struct NullFetcher;

#[cfg_attr(feature = "wasm", async_trait(?Send))]
#[cfg_attr(not(feature = "wasm"), async_trait)]
impl Fetcher for NullFetcher {
    async fn fetch(&self, _request: HttpRequest) -> Result<HttpResponse> {
        Err(anyhow!("Not found"))
    }
}
<commit_msg>Follow HTTP redirect header (#368)<commit_after>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://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.

#[cfg(feature = "wasm")]
pub mod js_fetcher;
pub mod mock_fetcher;

use crate::http::{HttpRequest, HttpResponse};
use crate::utils::{MaybeSend, MaybeSync};
use anyhow::{anyhow, Result};
use async_trait::async_trait;

/// An interface for fetching resources from network.
#[cfg_attr(feature = "wasm", async_trait(?Send))]
#[cfg_attr(not(feature = "wasm"), async_trait)]
pub trait Fetcher: MaybeSend + MaybeSync {
    async fn fetch(&self, request: HttpRequest) -> Result<HttpResponse>;
}

/// Uses `Get` method and returns response body,
/// iteratively following 301, 302, 303, 307, 308 redirection.
/// - Why this function is not put inside [`Fetcher`] trait?
///   If we declare `Fetcher::get` function with a default implementation,
///   we have to also add a constraint `where Self: Sized` to `Fetcher::get`,
///   because of https://github.com/rust-lang/rust/issues/51443 and in particular
///   https://docs.rs/async-trait/0.1.57/async_trait/#dyn-traits.
///   However, having such constraint `Self: Sized` prevent using `Fetcher::get` method on a
///   `dyn Fetcher` variable, because `dyn Fetcher` is not `Sized`.
pub async fn get(fetcher: &dyn Fetcher, url: impl ToString) -> Result<Vec<u8>> {
    let mut url = url.to_string();
    loop {
        let request = HttpRequest {
            body: vec![],
            headers: vec![],
            method: crate::http::Method::Get,
            url: url.to_string(),
        };
        let response = fetcher.fetch(request).await?;
        if matches!(response.status, 301 | 302 | 303 | 307 | 308) {
            let location = response.headers.into_iter().find_map(|(name, value)| {
                if name.eq_ignore_ascii_case(http::header::LOCATION.as_str()) {
                    Some(value)
                } else {
                    None
                }
            });
            if let Some(location) = location {
                url = location;
                continue;
            }
        }
        return Ok(response.body);
    }
}

pub const NULL_FETCHER: NullFetcher = NullFetcher {};

pub struct NullFetcher;

#[cfg_attr(feature = "wasm", async_trait(?Send))]
#[cfg_attr(not(feature = "wasm"), async_trait)]
impl Fetcher for NullFetcher {
    async fn fetch(&self, _request: HttpRequest) -> Result<HttpResponse> {
        Err(anyhow!("Not found"))
    }
}
<|endoftext|>


<commit_before>use cppn_ext::cppn::{Cppn, CppnNode};
use cppn_ext::position::{Position, Position3d};
use cppn_ext::activation_function::{GeometricActivationFunction, ActivationFunction};
use weight::{Weight, WeightRange, WeightPerturbanceMethod};
use substrate::{Substrate, SubstrateConfiguration, Node, NetworkBuilder};
use behavioral_bitvec::BehavioralBitvec;
use genome::Genome;
use nsga2::driver::Driver;
use nsga2::population::{UnratedPopulation, RatedPopulation};
use fitness::Fitness;
use rand::Rng;
use mating::{MatingMethod, MatingMethodWeights};
use prob::Prob;
use neuron::Neuron;
use std::marker::PhantomData;

pub type CppnGenome<AF> where AF: ActivationFunction = Genome<CppnNode<AF>>;

pub struct NeuronNetworkBuilder;

impl NetworkBuilder for NeuronNetworkBuilder {
    type POS = Position3d;
    type NT = Neuron;
    type G = ();

    fn new() -> Self {
        NeuronNetworkBuilder
    }

    fn add_node(&mut self, node: &Node<Self::POS, Self::NT>, param: f64) {
        unimplemented!()
    }

    fn add_link(&mut self,
                source_node: &Node<Self::POS, Self::NT>,
                target_node: &Node<Self::POS, Self::NT>,
                weight1: f64,
                weight2: f64) {
        unimplemented!()
    }
    fn graph(self) -> Self::G {
        ()
    }
}

const CPPN_OUTPUT_LINK_WEIGHT1: usize = 0;
const CPPN_OUTPUT_LINK_WEIGHT2: usize = 1;
const CPPN_OUTPUT_LINK_EXPRESSION: usize = 2;
const CPPN_OUTPUT_NODE_WEIGHT: usize = 3;

/// Develops a network out of the CPPN
///
/// Returns the BehavioralBitvec and Connection Cost of the developed network

fn develop_cppn<'a, P, AF, T, V>(cppn: &mut Cppn<CppnNode<AF>, Weight, ()>,
                                 substrate_config: &SubstrateConfiguration<'a, P, T>,
                                 visitor: &mut V,
                                 leo_threshold: f64)
                                 -> (BehavioralBitvec, f64)
    where P: Position,
          AF: ActivationFunction,
          V: NetworkBuilder<POS = P, NT = T>
{
    // our CPPN has four outputs: link weight 1, link weight 2, link expression output, node weight
    assert!(cppn.output_count() == 4);
    assert!(cppn.input_count() == 6);

    let nodes = substrate_config.nodes();
    let links = substrate_config.links();
    let null_position = substrate_config.null_position();

    let mut bitvec = BehavioralBitvec::new(4 * (nodes.len() + links.len()));
    let mut connection_cost = 0.0;

    // First visit all nodes

    for node in nodes.iter() {
        let inputs = [node.position.coords(), null_position.coords()];
        cppn.process(&inputs[..]);

        let link_weight1 = cppn.read_output(CPPN_OUTPUT_LINK_WEIGHT1).unwrap();
        let link_weight2 = cppn.read_output(CPPN_OUTPUT_LINK_WEIGHT2).unwrap();
        let link_expression = cppn.read_output(CPPN_OUTPUT_LINK_EXPRESSION).unwrap();
        let node_weight = cppn.read_output(CPPN_OUTPUT_NODE_WEIGHT).unwrap();

        bitvec.push(link_weight1);
        bitvec.push(link_weight2);
        bitvec.push(link_expression);
        bitvec.push(node_weight);

        visitor.add_node(node, node_weight)
    }

    for &(source_node, target_node) in links.iter() {
        let inputs = [source_node.position.coords(), target_node.position.coords()];
        cppn.process(&inputs[..]);

        let link_weight1 = cppn.read_output(CPPN_OUTPUT_LINK_WEIGHT1).unwrap();
        let link_weight2 = cppn.read_output(CPPN_OUTPUT_LINK_WEIGHT2).unwrap();
        let link_expression = cppn.read_output(CPPN_OUTPUT_LINK_EXPRESSION).unwrap();
        let node_weight = cppn.read_output(CPPN_OUTPUT_NODE_WEIGHT).unwrap();

        bitvec.push(link_weight1);
        bitvec.push(link_weight2);
        bitvec.push(link_expression);
        bitvec.push(node_weight);

        if link_expression > leo_threshold {
            let distance_sq = source_node.position.distance_square(&target_node.position);
            debug_assert!(distance_sq >= 0.0);
            connection_cost += distance_sq;
            visitor.add_link(source_node, target_node, link_weight1, link_weight2);
        }
    }

    return (bitvec, connection_cost);
}

/// Determines the domain fitness of `G`

pub trait DomainFitness<G>: Sync where G: Sync {
    fn fitness(&self, g: G) -> f32;
}

pub struct CppnDriver<'a, DOMFIT, G, P, T, NETBUILDER>
    where DOMFIT: DomainFitness<G> + 'a,
          P: Position + Sync + 'a,
          T: Sync + 'a,
          NETBUILDER: NetworkBuilder<POS = P, NT = T, G = G> + Sync
{
    mating_method_weights: MatingMethodWeights,
    activation_functions: Vec<GeometricActivationFunction>,
    mutate_element_prob: Prob,
    weight_perturbance: WeightPerturbanceMethod,
    link_weight_range: WeightRange,

    mutate_add_node_random_link_weight: bool,
    mutate_drop_node_tournament_k: usize,
    mutate_modify_node_tournament_k: usize,

    mate_retries: usize,

    link_expression_threshold: f64,

    substrate_configuration: SubstrateConfiguration<'a, P, T>,
    domain_fitness: &'a DOMFIT,
    _phantom: PhantomData<NETBUILDER>,
}

impl<'a, DOMFIT, G, P, T, NETBUILDER> CppnDriver<'a, DOMFIT, G, P, T, NETBUILDER>
    where DOMFIT: DomainFitness<G>,
          P: Position + Sync + 'a,
          T: Sync + 'a,
          NETBUILDER: NetworkBuilder<POS = P, NT = T, G = G> + Sync
{
    fn random_hidden_node<R>(&self, rng: &mut R) -> CppnNode<GeometricActivationFunction>
        where R: Rng
    {
        let af = *rng.choose(&self.activation_functions).unwrap();
        CppnNode::hidden(af)
    }
}

impl<'a, DOMFIT, G, P, T, NETBUILDER> Driver for CppnDriver<'a, DOMFIT, G, P, T, NETBUILDER>
    where DOMFIT: DomainFitness<G>,
          G: Sync,
          P: Position + Sync + 'a,
          T: Sync + 'a,
          NETBUILDER: NetworkBuilder<POS = P, NT = T, G = G> + Sync
{
    type IND = CppnGenome<GeometricActivationFunction>;
    type FIT = Fitness;

    /// Creates a random individual for use by the start generation.
    ///
    /// We start from a minimal topology.

    fn random_individual<R>(&self, rng: &mut R) -> Self::IND
        where R: Rng
    {
        let mut genome = Self::IND::new();

        // 6 inputs (x1,y1,z1, x2,y2,z2)
        let inp_x1 = genome.add_node(CppnNode::input(GeometricActivationFunction::Linear));
        let inp_y1 = genome.add_node(CppnNode::input(GeometricActivationFunction::Linear));
        let inp_z1 = genome.add_node(CppnNode::input(GeometricActivationFunction::Linear));
        let inp_x2 = genome.add_node(CppnNode::input(GeometricActivationFunction::Linear));
        let inp_y2 = genome.add_node(CppnNode::input(GeometricActivationFunction::Linear));
        let inp_z2 = genome.add_node(CppnNode::input(GeometricActivationFunction::Linear));

        // 1 bias node (constant input of 1.0)
        let bias = genome.add_node(CppnNode::bias(GeometricActivationFunction::Constant1));

        // 4 outputs (t,w,ex,r)
        let out_t = genome.add_node(CppnNode::output(GeometricActivationFunction::BipolarGaussian));
        let out_w = genome.add_node(CppnNode::output(GeometricActivationFunction::BipolarGaussian));
        let out_ex = genome.add_node(CppnNode::output(GeometricActivationFunction::Linear));
        let out_r = genome.add_node(CppnNode::output(GeometricActivationFunction::BipolarGaussian));

        // make those nodes above immutable for mutation and crossover, as we need them to
        // develop the CPPN.
        genome.protect_nodes();

        // XXX: Add x-distance and y-distance Gaussian seed nodes.
        // XXX: Add initial random connections?

        genome
    }

    fn fitness(&self, ind: &Self::IND) -> Self::FIT {
        let mut cppn = Cppn::new(ind.network());
        let mut net_builder = NETBUILDER::new();

        // let substrate: Substrate<Position3d, Neuron> = Substrate::new();
        // let cfg = substrate.to_configuration();

        let (behavioral_bitvec, connection_cost) = develop_cppn(&mut cppn,
                                                                &self.substrate_configuration,
                                                                &mut net_builder,
                                                                self.link_expression_threshold);

        // XXX: Domain specific fitness
        Fitness {
            domain_fitness: 0.0,
            behavioral_diversity: 0,
            connection_cost: connection_cost,
            behavioral_bitvec: behavioral_bitvec,
        }
    }

    fn mate<R>(&self, rng: &mut R, parent1: &Self::IND, parent2: &Self::IND) -> Self::IND
        where R: Rng
    {
        let mut offspring = parent1.clone();

        for _ in 0..self.mate_retries + 1 {
            let modified = match MatingMethod::random_with(&self.mating_method_weights, rng) {
                MatingMethod::MutateAddNode => {
                    let link_weight = if self.mutate_add_node_random_link_weight {
                        Some(self.link_weight_range.random_weight(rng))
                    } else {
                        // duplicate existing node weight
                        None
                    };
                    let hidden_node = self.random_hidden_node(rng);
                    offspring.mutate_add_node(hidden_node, link_weight, rng)
                }
                MatingMethod::MutateDropNode => {
                    offspring.mutate_drop_node(self.mutate_drop_node_tournament_k, rng)
                }
                MatingMethod::MutateModifyNode => {
                    let hidden_node = self.random_hidden_node(rng);
                    offspring.mutate_modify_node(hidden_node,
                                                 self.mutate_modify_node_tournament_k,
                                                 rng)
                }
                MatingMethod::MutateConnect => {
                    let link_weight = self.link_weight_range.random_weight(rng);
                    offspring.mutate_connect(link_weight, rng)
                }
                MatingMethod::MutateDisconnect => offspring.mutate_disconnect(rng),
                MatingMethod::MutateWeights => {
                    let modifications = offspring.mutate_weights(self.mutate_element_prob,
                                                                 &self.weight_perturbance,
                                                                 &self.link_weight_range,
                                                                 rng);
                    modifications != 0
                }
                MatingMethod::CrossoverWeights => {
                    let modifications = offspring.crossover_weights(parent2, rng);
                    modifications != 0
                }
            };

            if modified {
                break;
            }
        }

        warn!("mate(): Genome was NOT modified!");

        return offspring;
    }

    fn population_metric(&self, population: &mut RatedPopulation<Self::IND, Self::FIT>) {
        // Determine the behavioral_diversity as average hamming distance to all other individuals.
        // hamming distance is symmetric.
        let n = population.len();
        for i in 0..n {
            // determine  hehavioral diversity for `i`.
            let mut diversity_i = 0;

            // XXX: parallelize this loop
            for j in i + 1..n {
                let distance = population.fitness()[i]
                                   .behavioral_bitvec
                                   .hamming_distance(&population.fitness()[j].behavioral_bitvec);
                diversity_i += distance;
                population.fitness_mut()[j].behavioral_diversity += distance;
            }

            population.fitness_mut()[i].behavioral_diversity = diversity_i;
        }
    }
}
<commit_msg>Calculate domain specific fitness<commit_after>use cppn_ext::cppn::{Cppn, CppnNode};
use cppn_ext::position::{Position, Position3d};
use cppn_ext::activation_function::{GeometricActivationFunction, ActivationFunction};
use weight::{Weight, WeightRange, WeightPerturbanceMethod};
use substrate::{Substrate, SubstrateConfiguration, Node, NetworkBuilder};
use behavioral_bitvec::BehavioralBitvec;
use genome::Genome;
use nsga2::driver::Driver;
use nsga2::population::{UnratedPopulation, RatedPopulation};
use fitness::Fitness;
use rand::Rng;
use mating::{MatingMethod, MatingMethodWeights};
use prob::Prob;
use neuron::Neuron;
use std::marker::PhantomData;

pub type CppnGenome<AF> where AF: ActivationFunction = Genome<CppnNode<AF>>;

pub struct NeuronNetworkBuilder;

impl NetworkBuilder for NeuronNetworkBuilder {
    type POS = Position3d;
    type NT = Neuron;
    type G = ();

    fn new() -> Self {
        NeuronNetworkBuilder
    }

    fn add_node(&mut self, node: &Node<Self::POS, Self::NT>, param: f64) {
        unimplemented!()
    }

    fn add_link(&mut self,
                source_node: &Node<Self::POS, Self::NT>,
                target_node: &Node<Self::POS, Self::NT>,
                weight1: f64,
                weight2: f64) {
        unimplemented!()
    }
    fn graph(self) -> Self::G {
        ()
    }
}

const CPPN_OUTPUT_LINK_WEIGHT1: usize = 0;
const CPPN_OUTPUT_LINK_WEIGHT2: usize = 1;
const CPPN_OUTPUT_LINK_EXPRESSION: usize = 2;
const CPPN_OUTPUT_NODE_WEIGHT: usize = 3;

/// Develops a network out of the CPPN
///
/// Returns the BehavioralBitvec and Connection Cost of the developed network

fn develop_cppn<'a, P, AF, T, V>(cppn: &mut Cppn<CppnNode<AF>, Weight, ()>,
                                 substrate_config: &SubstrateConfiguration<'a, P, T>,
                                 visitor: &mut V,
                                 leo_threshold: f64)
                                 -> (BehavioralBitvec, f64)
    where P: Position,
          AF: ActivationFunction,
          V: NetworkBuilder<POS = P, NT = T>
{
    // our CPPN has four outputs: link weight 1, link weight 2, link expression output, node weight
    assert!(cppn.output_count() == 4);
    assert!(cppn.input_count() == 6);

    let nodes = substrate_config.nodes();
    let links = substrate_config.links();
    let null_position = substrate_config.null_position();

    let mut bitvec = BehavioralBitvec::new(4 * (nodes.len() + links.len()));
    let mut connection_cost = 0.0;

    // First visit all nodes

    for node in nodes.iter() {
        let inputs = [node.position.coords(), null_position.coords()];
        cppn.process(&inputs[..]);

        let link_weight1 = cppn.read_output(CPPN_OUTPUT_LINK_WEIGHT1).unwrap();
        let link_weight2 = cppn.read_output(CPPN_OUTPUT_LINK_WEIGHT2).unwrap();
        let link_expression = cppn.read_output(CPPN_OUTPUT_LINK_EXPRESSION).unwrap();
        let node_weight = cppn.read_output(CPPN_OUTPUT_NODE_WEIGHT).unwrap();

        bitvec.push(link_weight1);
        bitvec.push(link_weight2);
        bitvec.push(link_expression);
        bitvec.push(node_weight);

        visitor.add_node(node, node_weight)
    }

    for &(source_node, target_node) in links.iter() {
        let inputs = [source_node.position.coords(), target_node.position.coords()];
        cppn.process(&inputs[..]);

        let link_weight1 = cppn.read_output(CPPN_OUTPUT_LINK_WEIGHT1).unwrap();
        let link_weight2 = cppn.read_output(CPPN_OUTPUT_LINK_WEIGHT2).unwrap();
        let link_expression = cppn.read_output(CPPN_OUTPUT_LINK_EXPRESSION).unwrap();
        let node_weight = cppn.read_output(CPPN_OUTPUT_NODE_WEIGHT).unwrap();

        bitvec.push(link_weight1);
        bitvec.push(link_weight2);
        bitvec.push(link_expression);
        bitvec.push(node_weight);

        if link_expression > leo_threshold {
            let distance_sq = source_node.position.distance_square(&target_node.position);
            debug_assert!(distance_sq >= 0.0);
            connection_cost += distance_sq;
            visitor.add_link(source_node, target_node, link_weight1, link_weight2);
        }
    }

    return (bitvec, connection_cost);
}

/// Determines the domain fitness of `G`

pub trait DomainFitness<G>: Sync where G: Sync {
    fn fitness(&self, g: G) -> f32;
}

pub struct CppnDriver<'a, DOMFIT, G, P, T, NETBUILDER>
    where DOMFIT: DomainFitness<G> + 'a,
          P: Position + Sync + 'a,
          T: Sync + 'a,
          NETBUILDER: NetworkBuilder<POS = P, NT = T, G = G> + Sync
{
    mating_method_weights: MatingMethodWeights,
    activation_functions: Vec<GeometricActivationFunction>,
    mutate_element_prob: Prob,
    weight_perturbance: WeightPerturbanceMethod,
    link_weight_range: WeightRange,

    mutate_add_node_random_link_weight: bool,
    mutate_drop_node_tournament_k: usize,
    mutate_modify_node_tournament_k: usize,

    mate_retries: usize,

    link_expression_threshold: f64,

    substrate_configuration: SubstrateConfiguration<'a, P, T>,
    domain_fitness: &'a DOMFIT,
    _phantom: PhantomData<NETBUILDER>,
}

impl<'a, DOMFIT, G, P, T, NETBUILDER> CppnDriver<'a, DOMFIT, G, P, T, NETBUILDER>
    where DOMFIT: DomainFitness<G>,
          P: Position + Sync + 'a,
          T: Sync + 'a,
          NETBUILDER: NetworkBuilder<POS = P, NT = T, G = G> + Sync
{
    fn random_hidden_node<R>(&self, rng: &mut R) -> CppnNode<GeometricActivationFunction>
        where R: Rng
    {
        let af = *rng.choose(&self.activation_functions).unwrap();
        CppnNode::hidden(af)
    }
}

impl<'a, DOMFIT, G, P, T, NETBUILDER> Driver for CppnDriver<'a, DOMFIT, G, P, T, NETBUILDER>
    where DOMFIT: DomainFitness<G>,
          G: Sync,
          P: Position + Sync + 'a,
          T: Sync + 'a,
          NETBUILDER: NetworkBuilder<POS = P, NT = T, G = G> + Sync
{
    type IND = CppnGenome<GeometricActivationFunction>;
    type FIT = Fitness;

    /// Creates a random individual for use by the start generation.
    ///
    /// We start from a minimal topology.

    fn random_individual<R>(&self, rng: &mut R) -> Self::IND
        where R: Rng
    {
        let mut genome = Self::IND::new();

        // 6 inputs (x1,y1,z1, x2,y2,z2)
        let inp_x1 = genome.add_node(CppnNode::input(GeometricActivationFunction::Linear));
        let inp_y1 = genome.add_node(CppnNode::input(GeometricActivationFunction::Linear));
        let inp_z1 = genome.add_node(CppnNode::input(GeometricActivationFunction::Linear));
        let inp_x2 = genome.add_node(CppnNode::input(GeometricActivationFunction::Linear));
        let inp_y2 = genome.add_node(CppnNode::input(GeometricActivationFunction::Linear));
        let inp_z2 = genome.add_node(CppnNode::input(GeometricActivationFunction::Linear));

        // 1 bias node (constant input of 1.0)
        let bias = genome.add_node(CppnNode::bias(GeometricActivationFunction::Constant1));

        // 4 outputs (t,w,ex,r)
        let out_t = genome.add_node(CppnNode::output(GeometricActivationFunction::BipolarGaussian));
        let out_w = genome.add_node(CppnNode::output(GeometricActivationFunction::BipolarGaussian));
        let out_ex = genome.add_node(CppnNode::output(GeometricActivationFunction::Linear));
        let out_r = genome.add_node(CppnNode::output(GeometricActivationFunction::BipolarGaussian));

        // make those nodes above immutable for mutation and crossover, as we need them to
        // develop the CPPN.
        genome.protect_nodes();

        // XXX: Add x-distance and y-distance Gaussian seed nodes.
        // XXX: Add initial random connections?

        genome
    }

    fn fitness(&self, ind: &Self::IND) -> Self::FIT {
        let mut cppn = Cppn::new(ind.network());
        let mut net_builder = NETBUILDER::new();

        let (behavioral_bitvec, connection_cost) = develop_cppn(&mut cppn,
                                                                &self.substrate_configuration,
                                                                &mut net_builder,
                                                                self.link_expression_threshold);

        // Evaluate domain specific fitness
        let domain_fitness = self.domain_fitness.fitness(net_builder.graph());

        Fitness {
            domain_fitness: domain_fitness,
            behavioral_diversity: 0, // will be calculated in `population_metric`
            connection_cost: connection_cost,
            behavioral_bitvec: behavioral_bitvec,
        }
    }

    fn mate<R>(&self, rng: &mut R, parent1: &Self::IND, parent2: &Self::IND) -> Self::IND
        where R: Rng
    {
        let mut offspring = parent1.clone();

        for _ in 0..self.mate_retries + 1 {
            let modified = match MatingMethod::random_with(&self.mating_method_weights, rng) {
                MatingMethod::MutateAddNode => {
                    let link_weight = if self.mutate_add_node_random_link_weight {
                        Some(self.link_weight_range.random_weight(rng))
                    } else {
                        // duplicate existing node weight
                        None
                    };
                    let hidden_node = self.random_hidden_node(rng);
                    offspring.mutate_add_node(hidden_node, link_weight, rng)
                }
                MatingMethod::MutateDropNode => {
                    offspring.mutate_drop_node(self.mutate_drop_node_tournament_k, rng)
                }
                MatingMethod::MutateModifyNode => {
                    let hidden_node = self.random_hidden_node(rng);
                    offspring.mutate_modify_node(hidden_node,
                                                 self.mutate_modify_node_tournament_k,
                                                 rng)
                }
                MatingMethod::MutateConnect => {
                    let link_weight = self.link_weight_range.random_weight(rng);
                    offspring.mutate_connect(link_weight, rng)
                }
                MatingMethod::MutateDisconnect => offspring.mutate_disconnect(rng),
                MatingMethod::MutateWeights => {
                    let modifications = offspring.mutate_weights(self.mutate_element_prob,
                                                                 &self.weight_perturbance,
                                                                 &self.link_weight_range,
                                                                 rng);
                    modifications != 0
                }
                MatingMethod::CrossoverWeights => {
                    let modifications = offspring.crossover_weights(parent2, rng);
                    modifications != 0
                }
            };

            if modified {
                break;
            }
        }

        warn!("mate(): Genome was NOT modified!");

        return offspring;
    }

    fn population_metric(&self, population: &mut RatedPopulation<Self::IND, Self::FIT>) {
        // Determine the behavioral_diversity as average hamming distance to all other individuals.
        // hamming distance is symmetric.
        let n = population.len();
        for i in 0..n {
            // determine  hehavioral diversity for `i`.
            let mut diversity_i = 0;

            // XXX: parallelize this loop
            for j in i + 1..n {
                let distance = population.fitness()[i]
                                   .behavioral_bitvec
                                   .hamming_distance(&population.fitness()[j].behavioral_bitvec);
                diversity_i += distance;
                population.fitness_mut()[j].behavioral_diversity += distance;
            }

            population.fitness_mut()[i].behavioral_diversity = diversity_i;
        }
    }
}
<|endoftext|>


<commit_before>
use nom::{HexDisplay, IResult, Needed, ErrorKind};

struct Document {
    header: Header,
    body: Vec<Element>,
}

struct Header {}

#[derive(Debug)]
enum ElementData {
    Signed(i64),
    Unsigned(u64),
    Float(f64),
    PlainString(String),
    UTF8String(String),
    Date(u64),
    Master(Vec<Element>),
    Binary(Vec<u8>),
    Unknown(u64),
}

#[derive(Debug)]
pub struct Element {
    id: u64,
    size: u64,
    data: ElementData,
}

pub fn vint(input: &[u8]) -> IResult<&[u8], u64> {
    if input.len() == 0 {
        return IResult::Incomplete(Needed::Size(1));
    }

    let v = input[0];
    let len = v.leading_zeros();

    if len == 8 {
        return IResult::Error(ErrorKind::Custom(1));
    }

    if input.len() <= len as usize {
        return IResult::Incomplete(Needed::Size(1));
    }

    let mut val = (v ^ (1 << (7 - len))) as u64;

    // println!("vint {:08b} {:08b} {:08b} {}", val, v, (1 << (8 - len)), len);

    for i in 0..len as usize {
        val = (val << 8) | input[i + 1] as u64;
    }

    IResult::Done(&input[len as usize + 1..], val)
}

// The ID are represented in the specification as their binary representation
// do not drop the marker bit.
pub fn vid(input: &[u8]) -> IResult<&[u8], u64> {
    if input.len() == 0 {
        return IResult::Incomplete(Needed::Size(1));
    }

    let v = input[0];
    let len = v.leading_zeros();

    if len == 8 {
        return IResult::Error(ErrorKind::Custom(1));
    }

    if input.len() <= len as usize {
        return IResult::Incomplete(Needed::Size(1));
    }

    let mut val = v as u64;

    // println!("vint {:08b} {:08b} {:08b} {}", val, v, (1 << (8 - len)), len);

    for i in 0..len as usize {
        val = (val << 8) | input[i + 1] as u64;
    }

    IResult::Done(&input[len as usize + 1..], val)
}

fn parse_master(input: &[u8], size: u64) -> IResult<&[u8], ElementData> {
    map!(input,
         many0!(parse_element),
         |elem| ElementData::Master(elem))
}

fn parse_uint(input: &[u8], size: u64) -> IResult<&[u8], ElementData> {
    let mut val = 0;

    if size > 8 {
        return IResult::Error(ErrorKind::Custom(1));
    }

    for i in 0..size as usize {
        val = (val << 8) | input[i] as u64;
    }

    IResult::Done(&input[size as usize..], ElementData::Unsigned(val))
}

pub fn parse_uint_data(input: &[u8], size: u64) -> IResult<&[u8], u64> {
    let mut val = 0;

    if size > 8 {
        return IResult::Error(ErrorKind::Custom(1));
    }

    for i in 0..size as usize {
        val = (val << 8) | input[i] as u64;
    }

    IResult::Done(&input[size as usize..], val)
}

pub fn parse_int_data(input: &[u8], size: u64) -> IResult<&[u8], i64> {
    let mut val = 0;

    if size > 8 {
        return IResult::Error(ErrorKind::Custom(1));
    }

    for i in 0..size as usize {
        val = (val << 8) | input[i] as u64;
    }

    //FIXME: is that right?
    IResult::Done(&input[size as usize..], val as i64)
}

use std::str;

fn parse_str(input: &[u8], size: u64) -> IResult<&[u8], ElementData> {
    do_parse!(input,
        s: take_s!(size as usize) >>
        ( ElementData::PlainString(String::from_utf8(s.to_owned()).unwrap()) )
    )
}

pub fn parse_str_data(input: &[u8], size: u64) -> IResult<&[u8], String> {
    do_parse!(input,
        s: take_s!(size as usize) >>
        ( String::from_utf8(s.to_owned()).unwrap() ) //FIXME: maybe do not unwrap here
    )
}

pub fn parse_binary_data(input: &[u8], size: u64) -> IResult<&[u8], Vec<u8>> {
    do_parse!(input,
        s: take_s!(size as usize) >>
        ( s.to_owned() )
    )
}

//FIXME: handle default values
//FIXME: is that really following IEEE_754-1985 ?
pub fn parse_float_data(input: &[u8], size: u64) -> IResult<&[u8], f64> {
    use nom::{be_f32, be_f64};
    if size == 0 {
        IResult::Done(input, 0f64)
    } else if size == 4 {
        map!(input, flat_map!(take!(4), be_f32), |val| val as f64)
    } else if size == 8 {
        flat_map!(input, take!(8), be_f64)
    } else {
        IResult::Error(ErrorKind::Custom(1))
    }
}

fn parse_element_id(input: &[u8], id: u64, size: u64) -> IResult<&[u8], ElementData> {
    // println!("id: 0x{:X} size: {}", id, size);
    if input.len() < size as usize {
        return IResult::Incomplete(Needed::Size(size as usize));
    }

    match id {
        0x1A45DFA3 => parse_master(input, size),
        0x4286 => parse_uint(input, size),
        0x42F7 => parse_uint(input, size),
        0x42F2 => parse_uint(input, size),
        0x42F3 => parse_uint(input, size),
        0x4282 => parse_str(input, size),
        0x4287 => parse_uint(input, size),
        0x4285 => parse_uint(input, size),
        _ => IResult::Done(&input[size as usize..], ElementData::Unknown(id)),
    }
}

named!(pub parse_element<Element>,
    do_parse!(
        id : vid >>
        size: vint >>
        data: call!(parse_element_id, id, size) >>
        (Element { id, size, data })
    )
);

#[macro_export]
macro_rules! ebml_uint (
  ($i: expr, $id:expr) => ({
    use $crate::ebml::{vid, vint, parse_uint_data};
    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: apply!(parse_uint_data, size)
      >> (data)
    )
  })
);

#[macro_export]
macro_rules! ebml_int (
  ($i: expr, $id:expr) => ({
    use $crate::ebml::{vid, vint, parse_int_data};
    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: apply!(parse_int_data, size)
      >> (data)
    )
  })
);

#[macro_export]
macro_rules! ebml_float (
  ($i: expr, $id:expr) => ({
    use $crate::ebml::{vid, vint, parse_float_data};
    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: apply!(parse_float_data, size)
      >> (data)
    )
  })
);

#[macro_export]
macro_rules! ebml_str (
  ($i: expr, $id:expr) => ({
    use $crate::ebml::{vid, vint, parse_str_data};

    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: apply!(parse_str_data, size)
      >> (data)
    )
  })
);

#[macro_export]
macro_rules! ebml_binary (
  ($i: expr, $id:expr) => ({
    use $crate::ebml::{vid, vint, parse_binary_data};

    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: apply!(parse_binary_data, size)
      >> (data)
    )
  })
);

#[macro_export]
macro_rules! ebml_master (
  ($i: expr, $id:expr, $submac:ident!( $($args:tt)* )) => ({
    use $crate::ebml::{vid, vint};
    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: flat_map!(take!(size as usize), $submac!($($args)*))
      >> (data)
    )
  })
);

#[derive(Debug,Clone,PartialEq)]
pub struct EBMLHeader {
    pub version: u64,
    pub read_version: u64,
    pub max_id_length: u64,
    pub max_size_length: u64,
    pub doc_type: String,
    pub doc_type_version: u64,
    pub doc_type_read_version: u64,
}

named!(pub ebml_header<EBMLHeader>,
  ebml_master!(0x1A45DFA3,
    do_parse!(
      t: permutation!(
        ebml_uint!(0x4286), // version
        ebml_uint!(0x42F7), // read_version
        ebml_uint!(0x42F2), // max id length
        ebml_uint!(0x42F3), // max size length
        ebml_str!(0x4282),  // doctype
        ebml_uint!(0x4287), // doctype version
        ebml_uint!(0x4285)  // doctype_read version
      ) >>
      ({
        EBMLHeader {
          version:               t.0,
          read_version:          t.1,
          max_id_length:         t.2,
          max_size_length:       t.3,
          doc_type:              t.4,
          doc_type_version:      t.5,
          doc_type_read_version: t.6,

        }
      })
    )
  )
);

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

    const single_stream: &'static [u8] = include_bytes!("../assets/single_stream.mkv");

    #[test]
    fn variable_integer() {
        let val01 = [0b10000000];
        //        let val01 = [0b01000000, 0b1];

        match vint(&val01) {
            IResult::Done(_, v) => assert!(0 == v),
            _ => panic!(),
        }
    }

    /*
    #[test]
    fn header() {
        println!("{}", single_stream[..8].to_hex(8));
        println!("{:b} {:b}", single_stream[0], single_stream[1]);
        println!("{:#?}", parse_element(single_stream));
        panic!();
    }*/

    #[test]
    fn header() {
        println!("{}", single_stream[..8].to_hex(8));
        println!("{:b} {:b}", single_stream[0], single_stream[1]);
        println!("{:?}", ebml_header(&single_stream[..100]));
        panic!();
    }
}
<commit_msg>the skip_void parser can consume a Void element<commit_after>
use nom::{HexDisplay, IResult, Needed, ErrorKind};

struct Document {
    header: Header,
    body: Vec<Element>,
}

struct Header {}

#[derive(Debug)]
enum ElementData {
    Signed(i64),
    Unsigned(u64),
    Float(f64),
    PlainString(String),
    UTF8String(String),
    Date(u64),
    Master(Vec<Element>),
    Binary(Vec<u8>),
    Unknown(u64),
}

#[derive(Debug)]
pub struct Element {
    id: u64,
    size: u64,
    data: ElementData,
}

pub fn vint(input: &[u8]) -> IResult<&[u8], u64> {
    if input.len() == 0 {
        return IResult::Incomplete(Needed::Size(1));
    }

    let v = input[0];
    let len = v.leading_zeros();

    if len == 8 {
        return IResult::Error(ErrorKind::Custom(1));
    }

    if input.len() <= len as usize {
        return IResult::Incomplete(Needed::Size(1));
    }

    let mut val = (v ^ (1 << (7 - len))) as u64;

    // println!("vint {:08b} {:08b} {:08b} {}", val, v, (1 << (8 - len)), len);

    for i in 0..len as usize {
        val = (val << 8) | input[i + 1] as u64;
    }

    IResult::Done(&input[len as usize + 1..], val)
}

// The ID are represented in the specification as their binary representation
// do not drop the marker bit.
pub fn vid(input: &[u8]) -> IResult<&[u8], u64> {
    if input.len() == 0 {
        return IResult::Incomplete(Needed::Size(1));
    }

    let v = input[0];
    let len = v.leading_zeros();

    if len == 8 {
        return IResult::Error(ErrorKind::Custom(1));
    }

    if input.len() <= len as usize {
        return IResult::Incomplete(Needed::Size(1));
    }

    let mut val = v as u64;

    // println!("vint {:08b} {:08b} {:08b} {}", val, v, (1 << (8 - len)), len);

    for i in 0..len as usize {
        val = (val << 8) | input[i + 1] as u64;
    }

    IResult::Done(&input[len as usize + 1..], val)
}

fn parse_master(input: &[u8], size: u64) -> IResult<&[u8], ElementData> {
    map!(input,
         many0!(parse_element),
         |elem| ElementData::Master(elem))
}

fn parse_uint(input: &[u8], size: u64) -> IResult<&[u8], ElementData> {
    let mut val = 0;

    if size > 8 {
        return IResult::Error(ErrorKind::Custom(1));
    }

    for i in 0..size as usize {
        val = (val << 8) | input[i] as u64;
    }

    IResult::Done(&input[size as usize..], ElementData::Unsigned(val))
}

pub fn parse_uint_data(input: &[u8], size: u64) -> IResult<&[u8], u64> {
    let mut val = 0;

    if size > 8 {
        return IResult::Error(ErrorKind::Custom(1));
    }

    for i in 0..size as usize {
        val = (val << 8) | input[i] as u64;
    }

    IResult::Done(&input[size as usize..], val)
}

pub fn parse_int_data(input: &[u8], size: u64) -> IResult<&[u8], i64> {
    let mut val = 0;

    if size > 8 {
        return IResult::Error(ErrorKind::Custom(1));
    }

    for i in 0..size as usize {
        val = (val << 8) | input[i] as u64;
    }

    //FIXME: is that right?
    IResult::Done(&input[size as usize..], val as i64)
}

use std::str;

fn parse_str(input: &[u8], size: u64) -> IResult<&[u8], ElementData> {
    do_parse!(input,
        s: take_s!(size as usize) >>
        ( ElementData::PlainString(String::from_utf8(s.to_owned()).unwrap()) )
    )
}

pub fn parse_str_data(input: &[u8], size: u64) -> IResult<&[u8], String> {
    do_parse!(input,
        s: take_s!(size as usize) >>
        ( String::from_utf8(s.to_owned()).unwrap() ) //FIXME: maybe do not unwrap here
    )
}

pub fn parse_binary_data(input: &[u8], size: u64) -> IResult<&[u8], Vec<u8>> {
    do_parse!(input,
        s: take_s!(size as usize) >>
        ( s.to_owned() )
    )
}

//FIXME: handle default values
//FIXME: is that really following IEEE_754-1985 ?
pub fn parse_float_data(input: &[u8], size: u64) -> IResult<&[u8], f64> {
    use nom::{be_f32, be_f64};
    if size == 0 {
        IResult::Done(input, 0f64)
    } else if size == 4 {
        map!(input, flat_map!(take!(4), be_f32), |val| val as f64)
    } else if size == 8 {
        flat_map!(input, take!(8), be_f64)
    } else {
        IResult::Error(ErrorKind::Custom(1))
    }
}

fn parse_element_id(input: &[u8], id: u64, size: u64) -> IResult<&[u8], ElementData> {
    // println!("id: 0x{:X} size: {}", id, size);
    if input.len() < size as usize {
        return IResult::Incomplete(Needed::Size(size as usize));
    }

    match id {
        0x1A45DFA3 => parse_master(input, size),
        0x4286 => parse_uint(input, size),
        0x42F7 => parse_uint(input, size),
        0x42F2 => parse_uint(input, size),
        0x42F3 => parse_uint(input, size),
        0x4282 => parse_str(input, size),
        0x4287 => parse_uint(input, size),
        0x4285 => parse_uint(input, size),
        _ => IResult::Done(&input[size as usize..], ElementData::Unknown(id)),
    }
}

named!(pub parse_element<Element>,
    do_parse!(
        id : vid >>
        size: vint >>
        data: call!(parse_element_id, id, size) >>
        (Element { id, size, data })
    )
);

#[macro_export]
macro_rules! ebml_uint (
  ($i: expr, $id:expr) => ({
    use $crate::ebml::{vid, vint, parse_uint_data};
    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: apply!(parse_uint_data, size)
      >> (data)
    )
  })
);

#[macro_export]
macro_rules! ebml_int (
  ($i: expr, $id:expr) => ({
    use $crate::ebml::{vid, vint, parse_int_data};
    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: apply!(parse_int_data, size)
      >> (data)
    )
  })
);

#[macro_export]
macro_rules! ebml_float (
  ($i: expr, $id:expr) => ({
    use $crate::ebml::{vid, vint, parse_float_data};
    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: apply!(parse_float_data, size)
      >> (data)
    )
  })
);

#[macro_export]
macro_rules! ebml_str (
  ($i: expr, $id:expr) => ({
    use $crate::ebml::{vid, vint, parse_str_data};

    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: apply!(parse_str_data, size)
      >> (data)
    )
  })
);

#[macro_export]
macro_rules! ebml_binary (
  ($i: expr, $id:expr) => ({
    use $crate::ebml::{vid, vint, parse_binary_data};

    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: apply!(parse_binary_data, size)
      >> (data)
    )
  })
);

#[macro_export]
macro_rules! ebml_master (
  ($i: expr, $id:expr, $submac:ident!( $($args:tt)* )) => ({
    use $crate::ebml::{vid, vint};
    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: flat_map!(take!(size as usize), $submac!($($args)*))
      >> (data)
    )
  })
);

named!(pub skip_void,
do_parse!(
  id:   verify!(vid, |val:u64| val == 0xEC) >>
  size: vint >>
  data: take!(size) >>
  (data)
  ));

#[derive(Debug,Clone,PartialEq)]
pub struct EBMLHeader {
    pub version: u64,
    pub read_version: u64,
    pub max_id_length: u64,
    pub max_size_length: u64,
    pub doc_type: String,
    pub doc_type_version: u64,
    pub doc_type_read_version: u64,
}

named!(pub ebml_header<EBMLHeader>,
  ebml_master!(0x1A45DFA3,
    do_parse!(
      t: permutation_opt!(
        complete!(ebml_uint!(0x4286)), // version
        complete!(ebml_uint!(0x42F7)), // read_version
        complete!(ebml_uint!(0x42F2)), // max id length
        complete!(ebml_uint!(0x42F3)), // max size length
        complete!(ebml_str!(0x4282)),  // doctype
        complete!(ebml_uint!(0x4287)), // doctype version
        complete!(ebml_uint!(0x4285)),  // doctype_read version
        complete!(skip_void)?
      ) >>
      ({
        EBMLHeader {
          version:               t.0,
          read_version:          t.1,
          max_id_length:         t.2,
          max_size_length:       t.3,
          doc_type:              t.4,
          doc_type_version:      t.5,
          doc_type_read_version: t.6,

        }
      })
    )
  )
);

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

    const single_stream: &'static [u8] = include_bytes!("../assets/single_stream.mkv");
    const webm: &'static [u8] = include_bytes!("../assets/big-buck-bunny_trailer.webm");

    #[test]
    fn variable_integer() {
        let val01 = [0b10000000];
        //        let val01 = [0b01000000, 0b1];

        match vint(&val01) {
            IResult::Done(_, v) => assert!(0 == v),
            _ => panic!(),
        }
    }

    /*
    #[test]
    fn header() {
        println!("{}", single_stream[..8].to_hex(8));
        println!("{:b} {:b}", single_stream[0], single_stream[1]);
        println!("{:#?}", parse_element(single_stream));
        panic!();
    }*/

    #[test]
    fn mkv_header() {
        println!("{}", single_stream[..8].to_hex(8));
        println!("{:b} {:b}", single_stream[0], single_stream[1]);
        println!("{:?}", ebml_header(&single_stream[..100]));
        panic!();
    }

    #[test]
    fn webm_header() {
        println!("{}", webm[..8].to_hex(8));
        println!("{:?}", ebml_header(&webm[..100]));
        panic!();
    }
}
<|endoftext|>


<commit_before>/// Copyright (c) 2015, Takeru Ohta <phjgt308@gmail.com>
///
extern crate eetf;

use eetf::*;

#[test]
fn decode_empty() {
    let input = [];
    let expected = None;
    assert_eq!(expected, decode(&input));
}

#[test]
fn decode_small_integer() {
    let input = [131,97,5];
    let expected = 5;
    assert_eq!(Term::Int(expected), decode(&input).unwrap());
}

#[test]
fn decode_integer() {
    let input = [131,98,0,0,4,210];
    let expected = 1234;
    assert_eq!(Term::Int(expected), decode(&input).unwrap());

    let input = [131,98,255,255,251,46];
    let expected = -1234;
    assert_eq!(Term::Int(expected), decode(&input).unwrap());
}

// #[test]
// fn decode_float() {
//     let input = [131,99,49,46,50,51,51,57,57,57,57,57,57,57,57,57,57,57,57,57,56,53,55,57,101,43,48,49,0,0,0,0,0];
//     let expected = 12.34;
//     assert_eq!(Term::Float(expected), decode(&input).unwrap());
// }

// #[test]
// fn decode_new_float() {
//     let input = [131,70,64,40,174,20,122,225,71,174];
//     let expected = 12.34;
//     assert_eq!(Term::Float(expected), decode(&input).unwrap());
// }

#[test]
fn decode_atom() {
    let input = [131,100,0,4,104,111,103,101];
    let expected = "hoge".to_string();
    assert_eq!(Term::Atom(expected), decode(&input).unwrap());
}

#[test]
fn decode_nil() {
    let input = [131,106];
    let expected = vec![];
    assert_eq!(Term::List(expected), decode(&input).unwrap());
}
<commit_msg>Define assert_decode!: an auxiliary macro for unit-tests<commit_after>/// Copyright (c) 2015, Takeru Ohta <phjgt308@gmail.com>
///
extern crate eetf;

use eetf::*;

macro_rules! assert_decode {
    ($x:expr, $y:expr) => {
        assert_eq!($x, decode(&$y))
    }
}

#[test]
fn decode_empty() {
    let input = [];
    let expected = None;
    assert_decode!(expected, input);
}

#[test]
fn decode_small_integer() {
    let input = [131,97,5];
    let expected = 5;
    assert_decode!(Some(Term::Int(expected)), input);
}

#[test]
fn decode_integer() {
    let input = [131,98,0,0,4,210];
    let expected = 1234;
    assert_decode!(Some(Term::Int(expected)), input);

    let input = [131,98,255,255,251,46];
    let expected = -1234;
    assert_decode!(Some(Term::Int(expected)), input);
}

// #[test]
// fn decode_float() {
//     let input = [131,99,49,46,50,51,51,57,57,57,57,57,57,57,57,57,57,57,57,57,56,53,55,57,101,43,48,49,0,0,0,0,0];
//     let expected = 12.34;
//     assert_eq!(Term::Float(expected), decode(&input).unwrap());
// }

// #[test]
// fn decode_new_float() {
//     let input = [131,70,64,40,174,20,122,225,71,174];
//     let expected = 12.34;
//     assert_eq!(Term::Float(expected), decode(&input).unwrap());
// }

#[test]
fn decode_atom() {
    let input = [131,100,0,4,104,111,103,101];
    let expected = "hoge".to_string();
    assert_decode!(Some(Term::Atom(expected)), input);
}

#[test]
fn decode_nil() {
    let input = [131,106];
    let expected = vec![];
    assert_decode!(Some(Term::List(expected)), input);
}
<|endoftext|>


<commit_before>extern crate pihex;
use pihex::*;

const TEST_CASES: &[(u64, &str)] = &[(0, "243f"),
                                     (1, "43f6"),
                                     (2, "3f6a"),
                                     (3, "f6a8"),
                                     (4, "6a88"),
                                     (8, "85a3"),
                                     (12, "08d3"),
                                     (96, "c0ac"),
                                     (128, "9216"),
                                     (1024, "25d4"),
                                     (3704, "34c6"),
                                     (4096, "5a04"),
                                     (4527, "e4f3"),
                                     (8192, "77af"),
                                     (2090, "e674"),
                                     (2944, "5094"),
                                     (11027, "e70f"),
                                     (11742, "41fb"),
                                     (15367, "0c67")];

#[test]
fn bbp_pihex_test() {
    for &(d, hex) in TEST_CASES {
        assert_eq!(bbp::pihex(d), hex);
    }
}

#[test]
fn pihex_bellard_test() {
    for &(d, hex) in TEST_CASES {
        assert_eq!(bellard::pihex(d), hex);
    }
}

#[test]
fn bbp_pihex_bellard_pihex_test() {
    for i in 0..1000 {
        assert_eq!((i, bbp::pihex(i)), (i, bellard::pihex(i)));
    }
}
<commit_msg>reorder the test cases<commit_after>extern crate pihex;
use pihex::*;

const TEST_CASES: &[(u64, &str)] = &[(0, "243f"),
                                     (1, "43f6"),
                                     (2, "3f6a"),
                                     (3, "f6a8"),
                                     (4, "6a88"),
                                     (8, "85a3"),
                                     (12, "08d3"),
                                     (96, "c0ac"),
                                     (128, "9216"),
                                     (1024, "25d4"),
                                     (2090, "e674"),
                                     (2944, "5094"),
                                     (3704, "34c6"),
                                     (4096, "5a04"),
                                     (4527, "e4f3"),
                                     (8192, "77af"),
                                     (11027, "e70f"),
                                     (11742, "41fb"),
                                     (15367, "0c67")];

#[test]
fn bbp_pihex_test() {
    for &(d, hex) in TEST_CASES {
        assert_eq!(bbp::pihex(d), hex);
    }
}

#[test]
fn pihex_bellard_test() {
    for &(d, hex) in TEST_CASES {
        assert_eq!(bellard::pihex(d), hex);
    }
}

#[test]
fn bbp_pihex_bellard_pihex_test() {
    for i in 0..1000 {
        assert_eq!((i, bbp::pihex(i)), (i, bellard::pihex(i)));
    }
}
<|endoftext|>


<commit_before>#![cfg(unix)]
#[macro_use]
mod util;

use mio::net::{UnixDatagram, UnixListener, UnixStream};
use mio::unix::SocketAddr;
use mio::{Interests, Token};
use std::io::{self, IoSlice, IoSliceMut, Read, Write};
use std::net::Shutdown;
use std::sync::mpsc::{channel, Receiver};
use std::thread;
use tempdir::TempDir;
use util::{
    assert_send, assert_sync, assert_would_block, expect_events, expect_no_events, init_with_poll,
    ExpectEvent, TryRead, TryWrite,
};

const DATA1: &[u8] = b"Hello same host!";
const DATA2: &[u8] = b"Why hello mio!";

const DEFAULT_BUF_SIZE: usize = 64;
const DATA1_LEN: usize = 16;
const DATA2_LEN: usize = 14;

const LOCAL: Token = Token(0);
const LOCAL_CLONE: Token = Token(1);

#[test]
fn smoke_test() {
    let (mut poll, mut events) = init_with_poll();

    let (sync_sender, sync_receiver) = channel();
    let (handle, remote_addr) = echo_remote(1, sync_receiver);

    let path = remote_addr.as_pathname().expect("not a pathname");
    let mut local = assert_ok!(UnixStream::connect(path));
    assert_ok!(sync_sender.send(()));

    assert_ok!(poll.registry().register(
        &local,
        LOCAL,
        Interests::WRITABLE.add(Interests::READABLE)
    ));

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::WRITABLE)],
    );

    let mut buf = [0; DEFAULT_BUF_SIZE];
    assert_would_block(local.read(&mut buf));

    let wrote = assert_ok!(local.write(&DATA1));
    assert_eq!(wrote, DATA1_LEN);
    assert_ok!(local.flush());

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::READABLE)],
    );

    let read = assert_ok!(local.read(&mut buf));
    assert_eq!(read, DATA1_LEN);
    assert_eq!(&buf[..read], DATA1);
    assert_eq!(read, wrote, "unequal reads and writes");

    assert!(assert_ok!(local.take_error()).is_none());

    let bufs = [IoSlice::new(&DATA1), IoSlice::new(&DATA2)];
    let wrote = assert_ok!(local.write_vectored(&bufs));
    assert_eq!(wrote, DATA1_LEN + DATA2_LEN);

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::READABLE)],
    );

    let mut buf1 = [1; DATA1_LEN];
    let mut buf2 = [2; DATA2_LEN + 1];
    let mut bufs = [IoSliceMut::new(&mut buf1), IoSliceMut::new(&mut buf2)];
    let read = assert_ok!(local.read_vectored(&mut bufs));
    assert_eq!(read, DATA1_LEN + DATA2_LEN);
    assert_eq!(&buf1, DATA1);
    assert_eq!(&buf2[..DATA2.len()], DATA2);

    // Last byte should be unchanged
    assert_eq!(buf2[DATA2.len()], 2);

    // Close the connection to allow the remote to shutdown
    drop(local);
    assert_ok!(handle.join());
}

#[test]
fn is_send_and_sync() {
    assert_send::<UnixListener>();
    assert_sync::<UnixListener>();

    assert_send::<UnixStream>();
    assert_sync::<UnixStream>();

    assert_send::<UnixDatagram>();
    assert_sync::<UnixDatagram>();
}

#[test]
fn register() {
    let (mut poll, mut events) = init_with_poll();

    let (sync_sender, sync_receiver) = channel();
    let (handle, remote_addr) = echo_remote(1, sync_receiver);

    let path = remote_addr.as_pathname().expect("not a pathname");
    let local = assert_ok!(UnixStream::connect(path));
    assert_ok!(sync_sender.send(()));

    assert_ok!(poll.registry().register(&local, LOCAL, Interests::READABLE));

    expect_no_events(&mut poll, &mut events);

    // Close the connection to allow the remote to shutdown
    drop(local);
    assert_ok!(handle.join());
}

#[test]
fn reregister() {
    let (mut poll, mut events) = init_with_poll();

    let (sync_sender, sync_receiver) = channel();
    let (handle, remote_addr) = echo_remote(1, sync_receiver);

    let path = remote_addr.as_pathname().expect("not a pathname");
    let local = assert_ok!(UnixStream::connect(path));
    assert_ok!(sync_sender.send(()));

    assert_ok!(poll.registry().register(&local, LOCAL, Interests::READABLE));
    assert_ok!(poll
        .registry()
        .reregister(&local, LOCAL_CLONE, Interests::WRITABLE));

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL_CLONE, Interests::WRITABLE)],
    );

    // Close the connection to allow the remote to shutdown
    drop(local);
    assert_ok!(handle.join());
}

#[test]
fn deregister() {
    let (mut poll, mut events) = init_with_poll();

    let (sync_sender, sync_receiver) = channel();
    let (handle, remote_addr) = echo_remote(1, sync_receiver);

    let path = remote_addr.as_pathname().expect("not a pathname");
    let local = assert_ok!(UnixStream::connect(path));
    assert_ok!(sync_sender.send(()));

    assert_ok!(poll.registry().register(&local, LOCAL, Interests::WRITABLE));
    assert_ok!(poll.registry().deregister(&local));

    expect_no_events(&mut poll, &mut events);

    // Close the connection to allow the remote to shutdown
    drop(local);
    assert_ok!(handle.join());
}

#[test]
fn connect() {
    let (mut poll, mut events) = init_with_poll();

    let dir = assert_ok!(TempDir::new("uds"));
    let path = dir.path().join("foo");

    let remote = assert_ok!(UnixListener::bind(path.clone()));
    let local = assert_ok!(UnixStream::connect(path));

    let (sender_1, receiver_1) = channel();
    let (sender_2, receiver_2) = channel();

    let handle = thread::spawn(move || {
        let (stream, _) = assert_ok!(remote.accept());
        assert_ok!(receiver_1.recv());
        drop(stream);
        assert_ok!(sender_2.send(()));
    });

    assert_ok!(poll
        .registry()
        .register(&local, LOCAL, Interests::READABLE | Interests::WRITABLE));

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::WRITABLE)],
    );

    assert_ok!(sender_1.send(()));
    assert_ok!(receiver_2.recv());

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::READABLE)],
    );

    assert_ok!(handle.join());
}

#[test]
fn try_clone() {
    let (mut poll, mut events) = init_with_poll();

    let (sync_sender, sync_receiver) = channel();
    let (handle, remote_addr) = echo_remote(1, sync_receiver);

    let path = remote_addr.as_pathname().expect("not a pathname");
    let mut local_1 = assert_ok!(UnixStream::connect(path));
    assert_ok!(sync_sender.send(()));

    assert_ok!(poll
        .registry()
        .register(&local_1, LOCAL, Interests::WRITABLE));

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::WRITABLE)],
    );

    let mut buf = [0; DEFAULT_BUF_SIZE];
    let wrote = assert_ok!(local_1.write(&DATA1));
    assert_eq!(wrote, DATA1_LEN);

    let mut local_2 = assert_ok!(local_1.try_clone());

    // When using `try_clone` the `TcpStream` needs to be deregistered!
    assert_ok!(poll.registry().deregister(&local_1));
    drop(local_1);

    assert_ok!(poll
        .registry()
        .register(&local_2, LOCAL_CLONE, Interests::READABLE));

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL_CLONE, Interests::READABLE)],
    );

    let read = assert_ok!(local_2.read(&mut buf));
    assert_eq!(read, DATA1_LEN);
    assert_eq!(&buf[..read], DATA1);

    // Close the connection to allow the remote to shutdown
    drop(local_2);
    assert_ok!(handle.join());
}

#[test]
fn shutdown_read() {
    let (mut poll, mut events) = init_with_poll();

    let (sync_sender, sync_receiver) = channel();
    let (handle, remote_addr) = echo_remote(1, sync_receiver);

    let path = remote_addr.as_pathname().expect("not a pathname");
    let mut local = assert_ok!(UnixStream::connect(path));
    assert_ok!(sync_sender.send(()));

    assert_ok!(poll.registry().register(
        &local,
        LOCAL,
        Interests::WRITABLE.add(Interests::READABLE)
    ));

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::WRITABLE)],
    );

    let wrote = assert_ok!(local.write(DATA1));
    assert_eq!(wrote, DATA1_LEN);

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::READABLE)],
    );

    assert_ok!(local.shutdown(Shutdown::Read));

    // Shutting down the reading side is different on each platform. For example
    // on Linux based systems we can still read.
    #[cfg(any(
        target_os = "dragonfly",
        target_os = "freebsd",
        target_os = "ios",
        target_os = "macos",
        target_os = "netbsd",
        target_os = "openbsd"
    ))]
    {
        let mut buf = [0; DEFAULT_BUF_SIZE];
        let read = assert_ok!(local.read(&mut buf));
        assert_eq!(read, 0);
    }

    // Close the connection to allow the remote to shutdown
    drop(local);
    assert_ok!(handle.join());
}

#[test]
fn shutdown_write() {
    let (mut poll, mut events) = init_with_poll();

    let (sync_sender, sync_receiver) = channel();
    let (handle, remote_addr) = echo_remote(1, sync_receiver);

    let path = remote_addr.as_pathname().expect("not a pathname");
    let mut local = assert_ok!(UnixStream::connect(path));
    assert_ok!(sync_sender.send(()));

    assert_ok!(poll.registry().register(
        &local,
        LOCAL,
        Interests::WRITABLE.add(Interests::READABLE)
    ));

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::WRITABLE)],
    );

    let wrote = assert_ok!(local.write(DATA1));
    assert_eq!(wrote, DATA1_LEN);

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::READABLE)],
    );

    assert_ok!(local.shutdown(Shutdown::Write));

    let err = assert_err!(local.write(DATA2));
    assert_eq!(err.kind(), io::ErrorKind::BrokenPipe);

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::READABLE)],
    );

    // Read should be ok
    let mut buf = [0; DEFAULT_BUF_SIZE];
    let read = assert_ok!(local.read(&mut buf));
    assert_eq!(read, DATA1_LEN);
    assert_eq!(&buf[..read], DATA1);

    // Close the connection to allow the remote to shutdown
    drop(local);
    assert_ok!(handle.join());
}

#[test]
fn shutdown_both() {
    let (mut poll, mut events) = init_with_poll();

    let (sync_sender, sync_receiver) = channel();
    let (handle, remote_addr) = echo_remote(1, sync_receiver);

    let path = remote_addr.as_pathname().expect("not a pathname");
    let mut local = assert_ok!(UnixStream::connect(path));
    assert_ok!(sync_sender.send(()));

    assert_ok!(poll.registry().register(
        &local,
        LOCAL,
        Interests::WRITABLE.add(Interests::READABLE)
    ));

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::WRITABLE)],
    );

    let wrote = assert_ok!(local.write(DATA1));
    assert_eq!(wrote, DATA1_LEN);

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::READABLE)],
    );

    assert_ok!(local.shutdown(Shutdown::Both));

    // Shutting down the reading side is different on each platform. For example
    // on Linux based systems we can still read.
    #[cfg(any(
        target_os = "dragonfly",
        target_os = "freebsd",
        target_os = "ios",
        target_os = "macos",
        target_os = "netbsd",
        target_os = "openbsd"
    ))]
    {
        let mut buf = [0; DEFAULT_BUF_SIZE];
        let read = assert_ok!(local.read(&mut buf));
        assert_eq!(read, 0);
    }

    let err = assert_err!(local.write(DATA2));
    #[cfg(unix)]
    assert_eq!(err.kind(), io::ErrorKind::BrokenPipe);
    #[cfg(window)]
    assert_eq!(err.kind(), io::ErrorKind::ConnectionAbroted);

    // Close the connection to allow the remote to shutdown
    drop(local);
    assert_ok!(handle.join());
}

fn echo_remote(
    connections: usize,
    sync_receiver: Receiver<()>,
) -> (thread::JoinHandle<()>, SocketAddr) {
    let (addr_sender, addr_receiver) = channel();
    let handle = thread::spawn(move || {
        let dir = assert_ok!(TempDir::new("uds"));
        let path = dir.path().join("foo");
        let remote = assert_ok!(UnixListener::bind(path.clone()));
        let local_address = assert_ok!(remote.local_addr());
        assert_ok!(addr_sender.send(local_address));

        for _ in 0..connections {
            assert_ok!(sync_receiver.recv());
            let (mut local, _) = assert_ok!(remote.accept());

            // On Linux based system it will cause a connection reset
            // error when the reading side of the peer connection is
            // shutdown, we don't consider it an actual here.
            let (mut read, mut written) = (0, 0);
            let mut buf = [0; DEFAULT_BUF_SIZE];
            loop {
                let n = match local.try_read(&mut buf) {
                    Ok(Some(amount)) => {
                        read += amount;
                        amount
                    }
                    Ok(None) => continue,
                    Err(ref err) if err.kind() == io::ErrorKind::ConnectionReset => break,
                    Err(err) => panic!("{}", err),
                };
                match local.try_write(&buf[..n]) {
                    Ok(Some(amount)) => written += amount,
                    Ok(None) => continue,
                    Err(ref err) if err.kind() == io::ErrorKind::BrokenPipe => break,
                    Err(err) => panic!("{:?}", err),
                };
            }
            assert_eq!(read, written, "unequal reads and writes");
        }
    });
    (handle, assert_ok!(addr_receiver.recv()))
}
<commit_msg>Once the UDS echo listener reads 0, do not try to write<commit_after>#![cfg(unix)]
#[macro_use]
mod util;

use mio::net::{UnixDatagram, UnixListener, UnixStream};
use mio::unix::SocketAddr;
use mio::{Interests, Token};
use std::io::{self, IoSlice, IoSliceMut, Read, Write};
use std::net::Shutdown;
use std::sync::mpsc::{channel, Receiver};
use std::thread;
use tempdir::TempDir;
use util::{
    assert_send, assert_sync, assert_would_block, expect_events, expect_no_events, init_with_poll,
    ExpectEvent, TryRead, TryWrite,
};

const DATA1: &[u8] = b"Hello same host!";
const DATA2: &[u8] = b"Why hello mio!";

const DEFAULT_BUF_SIZE: usize = 64;
const DATA1_LEN: usize = 16;
const DATA2_LEN: usize = 14;

const LOCAL: Token = Token(0);
const LOCAL_CLONE: Token = Token(1);

#[test]
fn smoke_test() {
    let (mut poll, mut events) = init_with_poll();

    let (sync_sender, sync_receiver) = channel();
    let (handle, remote_addr) = echo_remote(1, sync_receiver);

    let path = remote_addr.as_pathname().expect("not a pathname");
    let mut local = assert_ok!(UnixStream::connect(path));
    assert_ok!(sync_sender.send(()));

    assert_ok!(poll.registry().register(
        &local,
        LOCAL,
        Interests::WRITABLE.add(Interests::READABLE)
    ));

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::WRITABLE)],
    );

    let mut buf = [0; DEFAULT_BUF_SIZE];
    assert_would_block(local.read(&mut buf));

    let wrote = assert_ok!(local.write(&DATA1));
    assert_eq!(wrote, DATA1_LEN);
    assert_ok!(local.flush());

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::READABLE)],
    );

    let read = assert_ok!(local.read(&mut buf));
    assert_eq!(read, DATA1_LEN);
    assert_eq!(&buf[..read], DATA1);
    assert_eq!(read, wrote, "unequal reads and writes");

    assert!(assert_ok!(local.take_error()).is_none());

    let bufs = [IoSlice::new(&DATA1), IoSlice::new(&DATA2)];
    let wrote = assert_ok!(local.write_vectored(&bufs));
    assert_eq!(wrote, DATA1_LEN + DATA2_LEN);

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::READABLE)],
    );

    let mut buf1 = [1; DATA1_LEN];
    let mut buf2 = [2; DATA2_LEN + 1];
    let mut bufs = [IoSliceMut::new(&mut buf1), IoSliceMut::new(&mut buf2)];
    let read = assert_ok!(local.read_vectored(&mut bufs));
    assert_eq!(read, DATA1_LEN + DATA2_LEN);
    assert_eq!(&buf1, DATA1);
    assert_eq!(&buf2[..DATA2.len()], DATA2);

    // Last byte should be unchanged
    assert_eq!(buf2[DATA2.len()], 2);

    // Close the connection to allow the remote to shutdown
    drop(local);
    assert_ok!(handle.join());
}

#[test]
fn is_send_and_sync() {
    assert_send::<UnixListener>();
    assert_sync::<UnixListener>();

    assert_send::<UnixStream>();
    assert_sync::<UnixStream>();

    assert_send::<UnixDatagram>();
    assert_sync::<UnixDatagram>();
}

#[test]
fn register() {
    let (mut poll, mut events) = init_with_poll();

    let (sync_sender, sync_receiver) = channel();
    let (handle, remote_addr) = echo_remote(1, sync_receiver);

    let path = remote_addr.as_pathname().expect("not a pathname");
    let local = assert_ok!(UnixStream::connect(path));
    assert_ok!(sync_sender.send(()));

    assert_ok!(poll.registry().register(&local, LOCAL, Interests::READABLE));

    expect_no_events(&mut poll, &mut events);

    // Close the connection to allow the remote to shutdown
    drop(local);
    assert_ok!(handle.join());
}

#[test]
fn reregister() {
    let (mut poll, mut events) = init_with_poll();

    let (sync_sender, sync_receiver) = channel();
    let (handle, remote_addr) = echo_remote(1, sync_receiver);

    let path = remote_addr.as_pathname().expect("not a pathname");
    let local = assert_ok!(UnixStream::connect(path));
    assert_ok!(sync_sender.send(()));

    assert_ok!(poll.registry().register(&local, LOCAL, Interests::READABLE));
    assert_ok!(poll
        .registry()
        .reregister(&local, LOCAL_CLONE, Interests::WRITABLE));

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL_CLONE, Interests::WRITABLE)],
    );

    // Close the connection to allow the remote to shutdown
    drop(local);
    assert_ok!(handle.join());
}

#[test]
fn deregister() {
    let (mut poll, mut events) = init_with_poll();

    let (sync_sender, sync_receiver) = channel();
    let (handle, remote_addr) = echo_remote(1, sync_receiver);

    let path = remote_addr.as_pathname().expect("not a pathname");
    let local = assert_ok!(UnixStream::connect(path));
    assert_ok!(sync_sender.send(()));

    assert_ok!(poll.registry().register(&local, LOCAL, Interests::WRITABLE));
    assert_ok!(poll.registry().deregister(&local));

    expect_no_events(&mut poll, &mut events);

    // Close the connection to allow the remote to shutdown
    drop(local);
    assert_ok!(handle.join());
}

#[test]
fn connect() {
    let (mut poll, mut events) = init_with_poll();

    let dir = assert_ok!(TempDir::new("uds"));
    let path = dir.path().join("foo");

    let remote = assert_ok!(UnixListener::bind(path.clone()));
    let local = assert_ok!(UnixStream::connect(path));

    let (sender_1, receiver_1) = channel();
    let (sender_2, receiver_2) = channel();

    let handle = thread::spawn(move || {
        let (stream, _) = assert_ok!(remote.accept());
        assert_ok!(receiver_1.recv());
        drop(stream);
        assert_ok!(sender_2.send(()));
    });

    assert_ok!(poll
        .registry()
        .register(&local, LOCAL, Interests::READABLE | Interests::WRITABLE));

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::WRITABLE)],
    );

    assert_ok!(sender_1.send(()));
    assert_ok!(receiver_2.recv());

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::READABLE)],
    );

    assert_ok!(handle.join());
}

#[test]
fn try_clone() {
    let (mut poll, mut events) = init_with_poll();

    let (sync_sender, sync_receiver) = channel();
    let (handle, remote_addr) = echo_remote(1, sync_receiver);

    let path = remote_addr.as_pathname().expect("not a pathname");
    let mut local_1 = assert_ok!(UnixStream::connect(path));
    assert_ok!(sync_sender.send(()));

    assert_ok!(poll
        .registry()
        .register(&local_1, LOCAL, Interests::WRITABLE));

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::WRITABLE)],
    );

    let mut buf = [0; DEFAULT_BUF_SIZE];
    let wrote = assert_ok!(local_1.write(&DATA1));
    assert_eq!(wrote, DATA1_LEN);

    let mut local_2 = assert_ok!(local_1.try_clone());

    // When using `try_clone` the `TcpStream` needs to be deregistered!
    assert_ok!(poll.registry().deregister(&local_1));
    drop(local_1);

    assert_ok!(poll
        .registry()
        .register(&local_2, LOCAL_CLONE, Interests::READABLE));

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL_CLONE, Interests::READABLE)],
    );

    let read = assert_ok!(local_2.read(&mut buf));
    assert_eq!(read, DATA1_LEN);
    assert_eq!(&buf[..read], DATA1);

    // Close the connection to allow the remote to shutdown
    drop(local_2);
    assert_ok!(handle.join());
}

#[test]
fn shutdown_read() {
    let (mut poll, mut events) = init_with_poll();

    let (sync_sender, sync_receiver) = channel();
    let (handle, remote_addr) = echo_remote(1, sync_receiver);

    let path = remote_addr.as_pathname().expect("not a pathname");
    let mut local = assert_ok!(UnixStream::connect(path));
    assert_ok!(sync_sender.send(()));

    assert_ok!(poll.registry().register(
        &local,
        LOCAL,
        Interests::WRITABLE.add(Interests::READABLE)
    ));

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::WRITABLE)],
    );

    let wrote = assert_ok!(local.write(DATA1));
    assert_eq!(wrote, DATA1_LEN);

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::READABLE)],
    );

    assert_ok!(local.shutdown(Shutdown::Read));

    // Shutting down the reading side is different on each platform. For example
    // on Linux based systems we can still read.
    #[cfg(any(
        target_os = "dragonfly",
        target_os = "freebsd",
        target_os = "ios",
        target_os = "macos",
        target_os = "netbsd",
        target_os = "openbsd"
    ))]
    {
        let mut buf = [0; DEFAULT_BUF_SIZE];
        let read = assert_ok!(local.read(&mut buf));
        assert_eq!(read, 0);
    }

    // Close the connection to allow the remote to shutdown
    drop(local);
    assert_ok!(handle.join());
}

#[test]
fn shutdown_write() {
    let (mut poll, mut events) = init_with_poll();

    let (sync_sender, sync_receiver) = channel();
    let (handle, remote_addr) = echo_remote(1, sync_receiver);

    let path = remote_addr.as_pathname().expect("not a pathname");
    let mut local = assert_ok!(UnixStream::connect(path));
    assert_ok!(sync_sender.send(()));

    assert_ok!(poll.registry().register(
        &local,
        LOCAL,
        Interests::WRITABLE.add(Interests::READABLE)
    ));

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::WRITABLE)],
    );

    let wrote = assert_ok!(local.write(DATA1));
    assert_eq!(wrote, DATA1_LEN);

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::READABLE)],
    );

    assert_ok!(local.shutdown(Shutdown::Write));

    let err = assert_err!(local.write(DATA2));
    assert_eq!(err.kind(), io::ErrorKind::BrokenPipe);

    // Read should be ok
    let mut buf = [0; DEFAULT_BUF_SIZE];
    let read = assert_ok!(local.read(&mut buf));
    assert_eq!(read, DATA1_LEN);
    assert_eq!(&buf[..read], DATA1);

    // Close the connection to allow the remote to shutdown
    drop(local);
    assert_ok!(handle.join());
}

#[test]
fn shutdown_both() {
    let (mut poll, mut events) = init_with_poll();

    let (sync_sender, sync_receiver) = channel();
    let (handle, remote_addr) = echo_remote(1, sync_receiver);

    let path = remote_addr.as_pathname().expect("not a pathname");
    let mut local = assert_ok!(UnixStream::connect(path));
    assert_ok!(sync_sender.send(()));

    assert_ok!(poll.registry().register(
        &local,
        LOCAL,
        Interests::WRITABLE.add(Interests::READABLE)
    ));

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::WRITABLE)],
    );

    let wrote = assert_ok!(local.write(DATA1));
    assert_eq!(wrote, DATA1_LEN);

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(LOCAL, Interests::READABLE)],
    );

    assert_ok!(local.shutdown(Shutdown::Both));

    // Shutting down the reading side is different on each platform. For example
    // on Linux based systems we can still read.
    #[cfg(any(
        target_os = "dragonfly",
        target_os = "freebsd",
        target_os = "ios",
        target_os = "macos",
        target_os = "netbsd",
        target_os = "openbsd"
    ))]
    {
        let mut buf = [0; DEFAULT_BUF_SIZE];
        let read = assert_ok!(local.read(&mut buf));
        assert_eq!(read, 0);
    }

    let err = assert_err!(local.write(DATA2));
    #[cfg(unix)]
    assert_eq!(err.kind(), io::ErrorKind::BrokenPipe);
    #[cfg(window)]
    assert_eq!(err.kind(), io::ErrorKind::ConnectionAbroted);

    // Close the connection to allow the remote to shutdown
    drop(local);
    assert_ok!(handle.join());
}

fn echo_remote(
    connections: usize,
    sync_receiver: Receiver<()>,
) -> (thread::JoinHandle<()>, SocketAddr) {
    let (addr_sender, addr_receiver) = channel();
    let handle = thread::spawn(move || {
        let dir = assert_ok!(TempDir::new("uds"));
        let path = dir.path().join("foo");
        let remote = assert_ok!(UnixListener::bind(path.clone()));
        let local_address = assert_ok!(remote.local_addr());
        assert_ok!(addr_sender.send(local_address));

        for _ in 0..connections {
            assert_ok!(sync_receiver.recv());
            let (mut local, _) = assert_ok!(remote.accept());

            // On Linux based system it will cause a connection reset
            // error when the reading side of the peer connection is
            // shutdown, we don't consider it an actual here.
            let (mut read, mut written) = (0, 0);
            let mut buf = [0; DEFAULT_BUF_SIZE];
            loop {
                let n = match local.try_read(&mut buf) {
                    Ok(Some(amount)) => {
                        read += amount;
                        amount
                    }
                    Ok(None) => continue,
                    Err(ref err) if err.kind() == io::ErrorKind::ConnectionReset => break,
                    Err(err) => panic!("{}", err),
                };
                if n == 0 {
                    break;
                }
                match local.try_write(&buf[..n]) {
                    Ok(Some(amount)) => written += amount,
                    Ok(None) => continue,
                    Err(ref err) if err.kind() == io::ErrorKind::BrokenPipe => break,
                    Err(err) => panic!("{:?}", err),
                };
            }
            assert_eq!(read, written, "unequal reads and writes");
        }
    });
    (handle, assert_ok!(addr_receiver.recv()))
}
<|endoftext|>