<commit_before>//! # Paging
//! Some code was borrowed from [Phil Opp's Blog](http://os.phil-opp.com/modifying-page-tables.html)

pub use self::entry::*;
use memory::{PAGE_SIZE, Frame, FrameAllocator};
use self::temporary_page::TemporaryPage;
pub use self::mapper::Mapper;
use core::ops::{Add, Deref, DerefMut};
use multiboot2::BootInformation;

pub mod entry;
mod mapper;
mod table;
mod temporary_page;

const ENTRY_COUNT: usize = 512;

// TODO move this to a separated module
const KERNEL_PERCPU_OFFSET: usize = 0xc000_0000;
/// Size of kernel percpu variables
const KERNEL_PERCPU_SIZE: usize = 64 * 1024; // 64kb

pub type PhysicalAddress = usize;
pub type VirtualAddress = usize;

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct Page {
    number: usize,
}

impl Page {
    pub fn containing_address(address: VirtualAddress) -> Page {
        assert!(address < 0x0000_8000_0000_0000 || address >= 0xffff_8000_0000_0000,
        "invalid address: 0x{:x}",
        address);
        Page { number: address / PAGE_SIZE }
    }

    pub fn start_address(&self) -> usize {
        self.number * PAGE_SIZE
    }

    fn p4_index(&self) -> usize {
        (self.number >> 27) & 0o777
    }
    fn p3_index(&self) -> usize {
        (self.number >> 18) & 0o777
    }
    fn p2_index(&self) -> usize {
        (self.number >> 9) & 0o777
    }
    fn p1_index(&self) -> usize {
        (self.number >> 0) & 0o777
    }

    pub fn range_inclusive(start: Page, end: Page) -> PageIter {
        PageIter {
            start: start,
            end: end,
        }
    }
}

impl Add<usize> for Page {
    type Output = Page;

    fn add(self, rhs: usize) -> Page {
        Page { number: self.number + rhs }
    }
}

#[derive(Clone)]
pub struct PageIter {
    start: Page,
    end: Page,
}

impl Iterator for PageIter {
    type Item = Page;

    fn next(&mut self) -> Option<Page> {
        if self.start <= self.end {
            let page = self.start;
            self.start.number += 1;
            Some(page)
        } else {
            None
        }
    }
}

pub struct ActivePageTable {
    mapper: Mapper,
}

impl Deref for ActivePageTable {
    type Target = Mapper;

    fn deref(&self) -> &Mapper {
        &self.mapper
    }
}

impl DerefMut for ActivePageTable {
    fn deref_mut(&mut self) -> &mut Mapper {
        &mut self.mapper
    }
}

impl ActivePageTable {
    pub unsafe fn new() -> ActivePageTable {
        ActivePageTable { mapper: Mapper::new() }
    }

    pub fn with<F>(&mut self,
                   table: &mut InactivePageTable,
                   temporary_page: &mut temporary_page::TemporaryPage,
                   f: F)
        where F: FnOnce(&mut Mapper)
    {
        use x86_64::registers::control_regs;

        {
            // backup the original CR3 address in order to restore it lately
            let backup = Frame::containing_address(control_regs::cr3().0 as usize);

            // map temporary_page to current p4 table
            let p4_table = temporary_page.map_table_frame(backup.clone(), self);

            // overwrite recursive mapping and flush the tlb to ensure the correct page translation
            self.p4_mut()[511].set(table.p4_frame.clone(), PRESENT | WRITABLE);

            // flush TLB
            self.flush_all();

            // execute f in the new context
            f(self);

            // restore recursive mapping to original p4 table and flush the tlb again
            p4_table[511].set(backup, PRESENT | WRITABLE);

            // flush TLB
            self.flush_all();
        }

        temporary_page.unmap(self);
    }

    /// Switch context
    pub fn switch(&mut self, new_table: InactivePageTable) -> InactivePageTable {
        use x86_64::PhysicalAddress;
        use x86_64::registers::control_regs;

        // store the ond table
        let old_table = InactivePageTable {
            p4_frame: Frame::containing_address(control_regs::cr3().0 as usize),
        };

        // switch to the new page table
        unsafe {
            control_regs::cr3_write(PhysicalAddress(new_table.p4_frame.start_address() as u64));
        }

        // return the old table
        old_table
    }

    /// Flush all the TLB table
    pub fn flush_all(&mut self) {
        use x86_64::instructions::tlb;

        unsafe { tlb::flush_all(); }
    }

    /// Get the CR3 address.
    pub unsafe fn address(&self) -> usize {
        use x86_64::registers::control_regs;
        control_regs::cr3().0 as usize
    }
}

pub struct InactivePageTable {
    p4_frame: Frame,
}

impl InactivePageTable {
    pub fn new(frame: Frame,
               active_table: &mut ActivePageTable,
               temporary_page: &mut TemporaryPage)
               -> InactivePageTable
    {
        {
            let table = temporary_page.map_table_frame(frame.clone(), active_table);

            // zero the page table
            table.zero();

            // set up recursive mapping for the table
            table[511].set(frame.clone(), PRESENT | WRITABLE);
        }
        temporary_page.unmap(active_table);

        InactivePageTable { p4_frame: frame }
    }
}

/// Remap the kernel
pub unsafe fn remap_the_kernel<A>(allocator: &mut A, boot_info: &BootInformation) -> ActivePageTable
    where A: FrameAllocator
{
    use core::ops::Range;

    // get the external data segments
    extern {
        /// The starting byte of the text (code) data segment.
        static mut __text_start: u8;
        /// The ending byte of the text (code) data segment.
        static mut __text_end: u8;
        /// The starting byte of the _.rodata_ (read-only data) segment.
        static mut __rodata_start: u8;
        /// The ending byte of the _.rodata_ (read-only data) segment.
        static mut __rodata_end: u8;
        /// The starting byte of the _.data_ segment.
        static mut __data_start: u8;
        /// The ending byte of the _.data_ segment.
        static mut __data_end: u8;
        /// The starting byte of the thread data segment
        static mut __tdata_start: u8;
        /// The ending byte of the thread data segment
        static mut __tdata_end: u8;
        /// The starting byte of the thread BSS segment
        static mut __tbss_start: u8;
        /// The ending byte of the thread BSS segment
        static mut __tbss_end: u8;
        /// The starting byte of the _.bss_ (uninitialized data) segment.
        static mut __bss_start: u8;
        /// The ending byte of the _.bss_ (uninitialized data) segment.
        static mut __bss_end: u8;
    }

    let mut temporary_page = TemporaryPage::new(Page { number: 0xcafebabe }, allocator);

    let mut active_table = unsafe { ActivePageTable::new() };
    let mut new_table = {
        let frame = allocator.allocate_frame().expect("no more frames");
        InactivePageTable::new(frame, &mut active_table, &mut temporary_page)
    };

    // TODO remap the kernel sections individual, set permissions and deal with the tbss section properly

    active_table.with(&mut new_table, &mut temporary_page, |mapper| {
        // Map tdata and tbss
        {
            // get the thread segment size.
            let size = & __tbss_end as *const _ as usize - & __tdata_start as *const _ as usize;

            // TODO add support to multiple CPUs (replace the numeric constant)
            let start = KERNEL_PERCPU_OFFSET + KERNEL_PERCPU_SIZE * 1;
            let end = start + size;

            let start_page = Page::containing_address(start as VirtualAddress);
            let end_page = Page::containing_address((end - 1) as VirtualAddress);
            for page in Page::range_inclusive(start_page, end_page) {
                mapper.map(page, PRESENT | GLOBAL | NO_EXECUTE | WRITABLE, allocator);
            }
        }

        // Function to remap the kernel sections individually. We can no longer use the multiboot
        // sections because of the thread base sections.
        let mut remap = move |start: usize, end: usize, flags: EntryFlags, mapper: &mut Mapper, allocator: &mut A| {
            assert!(start as usize % PAGE_SIZE == 0, "sections need to be page aligned");

            let start_frame = Frame::containing_address(start);
            let end_frame = Frame::containing_address(end - 1);
            for frame in Frame::range_inclusive(start_frame, end_frame) {
                mapper.identity_map(frame, flags, allocator);
            }
        };

        // Remap the kernel with `flags`
        let mut remap_section = move |start: &u8, end: &u8, flags: EntryFlags, mapper: &mut Mapper, allocator: &mut A| {
            remap(start as *const _ as usize, end as *const _ as usize, flags, mapper, allocator);
        };

        // Remap text read-only
        remap_section(& __text_start, & __text_end, PRESENT | GLOBAL, mapper, allocator);
        // Remap rodata read-only, no execute
        remap_section(& __rodata_start, & __rodata_end, PRESENT | GLOBAL | NO_EXECUTE, mapper, allocator);
        // Remap data writable, no execute
        remap_section(& __data_start, & __data_end, PRESENT | GLOBAL | NO_EXECUTE | WRITABLE, mapper, allocator);
        // Remap tdata master writable, no execute
        remap_section(& __tdata_start, & __tdata_end, PRESENT | GLOBAL | NO_EXECUTE, mapper, allocator);
        // Remap bss writable, no execute
        remap_section(& __bss_start, & __bss_end, PRESENT | GLOBAL | NO_EXECUTE | WRITABLE, mapper, allocator);

        // identity map the VGA text buffer
        let vga_buffer_frame = Frame::containing_address(0xb8000);
        mapper.identity_map(vga_buffer_frame, WRITABLE, allocator);

        // identity map the multiboot info structure
        let multiboot_start = Frame::containing_address(boot_info.start_address());
        let multiboot_end = Frame::containing_address(boot_info.end_address() - 1);
        for frame in Frame::range_inclusive(multiboot_start, multiboot_end) {
            mapper.identity_map(frame, PRESENT, allocator);
        }
    });

    // switch context to start using the new page table
    let old_table = active_table.switch(new_table);

    // turn the old p4 page into a guard page
    let old_p4_page = Page::containing_address(old_table.p4_frame.start_address());
    active_table.unmap(old_p4_page, allocator);
    println!("guard page at {:#x}", old_p4_page.start_address());

    active_table
}
<commit_msg>Initialize tcb<commit_after>//! # Paging
//! Some code was borrowed from [Phil Opp's Blog](http://os.phil-opp.com/modifying-page-tables.html)

use core::{mem, ptr};
use core::ops::{Add, Deref, DerefMut};
use memory::{PAGE_SIZE, Frame, FrameAllocator};
use self::temporary_page::TemporaryPage;
pub use self::mapper::Mapper;
pub use self::entry::*;
use multiboot2::BootInformation;

pub mod entry;
mod mapper;
mod table;
mod temporary_page;

const ENTRY_COUNT: usize = 512;

// TODO move this to a separated module
const KERNEL_PERCPU_OFFSET: usize = 0xc000_0000;
/// Size of kernel percpu variables
const KERNEL_PERCPU_SIZE: usize = 64 * 1024; // 64kb

pub type PhysicalAddress = usize;
pub type VirtualAddress = usize;

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct Page {
    number: usize,
}

impl Page {
    pub fn containing_address(address: VirtualAddress) -> Page {
        assert!(address < 0x0000_8000_0000_0000 || address >= 0xffff_8000_0000_0000,
        "invalid address: 0x{:x}",
        address);
        Page { number: address / PAGE_SIZE }
    }

    pub fn start_address(&self) -> usize {
        self.number * PAGE_SIZE
    }

    fn p4_index(&self) -> usize {
        (self.number >> 27) & 0o777
    }
    fn p3_index(&self) -> usize {
        (self.number >> 18) & 0o777
    }
    fn p2_index(&self) -> usize {
        (self.number >> 9) & 0o777
    }
    fn p1_index(&self) -> usize {
        (self.number >> 0) & 0o777
    }

    pub fn range_inclusive(start: Page, end: Page) -> PageIter {
        PageIter {
            start: start,
            end: end,
        }
    }
}

impl Add<usize> for Page {
    type Output = Page;

    fn add(self, rhs: usize) -> Page {
        Page { number: self.number + rhs }
    }
}

#[derive(Clone)]
pub struct PageIter {
    start: Page,
    end: Page,
}

impl Iterator for PageIter {
    type Item = Page;

    fn next(&mut self) -> Option<Page> {
        if self.start <= self.end {
            let page = self.start;
            self.start.number += 1;
            Some(page)
        } else {
            None
        }
    }
}

pub struct ActivePageTable {
    mapper: Mapper,
}

impl Deref for ActivePageTable {
    type Target = Mapper;

    fn deref(&self) -> &Mapper {
        &self.mapper
    }
}

impl DerefMut for ActivePageTable {
    fn deref_mut(&mut self) -> &mut Mapper {
        &mut self.mapper
    }
}

impl ActivePageTable {
    pub unsafe fn new() -> ActivePageTable {
        ActivePageTable { mapper: Mapper::new() }
    }

    pub fn with<F>(&mut self,
                   table: &mut InactivePageTable,
                   temporary_page: &mut temporary_page::TemporaryPage,
                   f: F)
        where F: FnOnce(&mut Mapper)
    {
        use x86_64::registers::control_regs;

        {
            // backup the original CR3 address in order to restore it lately
            let backup = Frame::containing_address(control_regs::cr3().0 as usize);

            // map temporary_page to current p4 table
            let p4_table = temporary_page.map_table_frame(backup.clone(), self);

            // overwrite recursive mapping and flush the tlb to ensure the correct page translation
            self.p4_mut()[511].set(table.p4_frame.clone(), PRESENT | WRITABLE);

            // flush TLB
            self.flush_all();

            // execute f in the new context
            f(self);

            // restore recursive mapping to original p4 table and flush the tlb again
            p4_table[511].set(backup, PRESENT | WRITABLE);

            // flush TLB
            self.flush_all();
        }

        temporary_page.unmap(self);
    }

    /// Switch context
    pub fn switch(&mut self, new_table: InactivePageTable) -> InactivePageTable {
        use x86_64::PhysicalAddress;
        use x86_64::registers::control_regs;

        // store the ond table
        let old_table = InactivePageTable {
            p4_frame: Frame::containing_address(control_regs::cr3().0 as usize),
        };

        // switch to the new page table
        unsafe {
            control_regs::cr3_write(PhysicalAddress(new_table.p4_frame.start_address() as u64));
        }

        // return the old table
        old_table
    }

    /// Flush all the TLB table
    pub fn flush_all(&mut self) {
        use x86_64::instructions::tlb;

        unsafe { tlb::flush_all(); }
    }

    /// Get the CR3 address.
    pub unsafe fn address(&self) -> usize {
        use x86_64::registers::control_regs;
        control_regs::cr3().0 as usize
    }
}

pub struct InactivePageTable {
    p4_frame: Frame,
}

impl InactivePageTable {
    pub fn new(frame: Frame,
               active_table: &mut ActivePageTable,
               temporary_page: &mut TemporaryPage)
               -> InactivePageTable
    {
        {
            let table = temporary_page.map_table_frame(frame.clone(), active_table);

            // zero the page table
            table.zero();

            // set up recursive mapping for the table
            table[511].set(frame.clone(), PRESENT | WRITABLE);
        }
        temporary_page.unmap(active_table);

        InactivePageTable { p4_frame: frame }
    }
}

/// Copy tdata, clear tbss, set TCB self pointer
unsafe fn init_tcb(cpu_id: usize) -> usize {
    extern {
        /// The starting byte of the thread data segment
        static mut __tdata_start: u8;
        /// The ending byte of the thread data segment
        static mut __tdata_end: u8;
        /// The starting byte of the thread BSS segment
        static mut __tbss_start: u8;
        /// The ending byte of the thread BSS segment
        static mut __tbss_end: u8;
    }

    let tcb_offset;
    {
        let size = & __tbss_end as *const _ as usize - & __tdata_start as *const _ as usize;
        let tbss_offset = & __tbss_start as *const _ as usize - & __tdata_start as *const _ as usize;

        let start = KERNEL_PERCPU_OFFSET + KERNEL_PERCPU_SIZE * cpu_id;
        let end = start + size;
        tcb_offset = end - mem::size_of::<usize>();

        // copy data
        ptr::copy(& __tdata_start as *const u8, start as *mut u8, tbss_offset);

        // zero .tbss
        ptr::write_bytes((start + tbss_offset) as *mut u8, 0, size - tbss_offset);

        *(tcb_offset as *mut usize) = end;
    }

    tcb_offset
}

/// Remap the kernel
pub unsafe fn remap_the_kernel<A>(allocator: &mut A, boot_info: &BootInformation) -> ActivePageTable
    where A: FrameAllocator
{
    use core::ops::Range;

    // get the external data segments
    extern {
        /// The starting byte of the text (code) data segment.
        static mut __text_start: u8;
        /// The ending byte of the text (code) data segment.
        static mut __text_end: u8;
        /// The starting byte of the _.rodata_ (read-only data) segment.
        static mut __rodata_start: u8;
        /// The ending byte of the _.rodata_ (read-only data) segment.
        static mut __rodata_end: u8;
        /// The starting byte of the _.data_ segment.
        static mut __data_start: u8;
        /// The ending byte of the _.data_ segment.
        static mut __data_end: u8;
        /// The starting byte of the thread data segment
        static mut __tdata_start: u8;
        /// The ending byte of the thread data segment
        static mut __tdata_end: u8;
        /// The starting byte of the thread BSS segment
        static mut __tbss_start: u8;
        /// The ending byte of the thread BSS segment
        static mut __tbss_end: u8;
        /// The starting byte of the _.bss_ (uninitialized data) segment.
        static mut __bss_start: u8;
        /// The ending byte of the _.bss_ (uninitialized data) segment.
        static mut __bss_end: u8;
    }

    let mut temporary_page = TemporaryPage::new(Page { number: 0xcafebabe }, allocator);

    let mut active_table = unsafe { ActivePageTable::new() };
    let mut new_table = {
        let frame = allocator.allocate_frame().expect("no more frames");
        InactivePageTable::new(frame, &mut active_table, &mut temporary_page)
    };

    // TODO remap the kernel sections individual, set permissions and deal with the tbss section properly

    active_table.with(&mut new_table, &mut temporary_page, |mapper| {
        // Map tdata and tbss
        {
            // get the thread segment size.
            let size = & __tbss_end as *const _ as usize - & __tdata_start as *const _ as usize;

            // TODO add support to multiple CPUs (replace the numeric constant)
            let start = KERNEL_PERCPU_OFFSET + KERNEL_PERCPU_SIZE * 1;
            let end = start + size;

            let start_page = Page::containing_address(start as VirtualAddress);
            let end_page = Page::containing_address((end - 1) as VirtualAddress);
            for page in Page::range_inclusive(start_page, end_page) {
                mapper.map(page, PRESENT | GLOBAL | NO_EXECUTE | WRITABLE, allocator);
            }
        }

        // Function to remap the kernel sections individually. We can no longer use the multiboot
        // sections because of the thread base sections.
        let mut remap = move |start: usize, end: usize, flags: EntryFlags, mapper: &mut Mapper, allocator: &mut A| {
            assert!(start as usize % PAGE_SIZE == 0, "sections need to be page aligned");

            let start_frame = Frame::containing_address(start);
            let end_frame = Frame::containing_address(end - 1);
            for frame in Frame::range_inclusive(start_frame, end_frame) {
                mapper.identity_map(frame, flags, allocator);
            }
        };

        // Remap the kernel with `flags`
        let mut remap_section = move |start: &u8, end: &u8, flags: EntryFlags, mapper: &mut Mapper, allocator: &mut A| {
            remap(start as *const _ as usize, end as *const _ as usize, flags, mapper, allocator);
        };

        // Remap text read-only
        remap_section(& __text_start, & __text_end, PRESENT | GLOBAL, mapper, allocator);
        // Remap rodata read-only, no execute
        remap_section(& __rodata_start, & __rodata_end, PRESENT | GLOBAL | NO_EXECUTE, mapper, allocator);
        // Remap data writable, no execute
        remap_section(& __data_start, & __data_end, PRESENT | GLOBAL | NO_EXECUTE | WRITABLE, mapper, allocator);
        // Remap tdata master writable, no execute
        remap_section(& __tdata_start, & __tdata_end, PRESENT | GLOBAL | NO_EXECUTE, mapper, allocator);
        // Remap bss writable, no execute
        remap_section(& __bss_start, & __bss_end, PRESENT | GLOBAL | NO_EXECUTE | WRITABLE, mapper, allocator);

        // identity map the VGA text buffer
        let vga_buffer_frame = Frame::containing_address(0xb8000);
        mapper.identity_map(vga_buffer_frame, WRITABLE, allocator);

        // identity map the multiboot info structure
        let multiboot_start = Frame::containing_address(boot_info.start_address());
        let multiboot_end = Frame::containing_address(boot_info.end_address() - 1);
        for frame in Frame::range_inclusive(multiboot_start, multiboot_end) {
            mapper.identity_map(frame, PRESENT, allocator);
        }
    });

    // switch context to start using the new page table
    let old_table = active_table.switch(new_table);

    // turn the old p4 page into a guard page
    let old_p4_page = Page::containing_address(old_table.p4_frame.start_address());
    active_table.unmap(old_p4_page, allocator);
    println!("guard page at {:#x}", old_p4_page.start_address());

    // initialize tcb
    // TODO make the CPU id dynamic
    init_tcb(1);

    active_table
}
<|endoftext|>


<commit_before>use dslite_parser;
use header_parser;
use inflector::Inflector;
use svd;

trait StringEx {
    fn fix_name(&self) -> String;
}

impl StringEx for String {
    fn fix_name(&self) -> String {
        self.to_screaming_snake_case()
            // Some fixups to make the names look nicer
            .replace("RTC_REAL_TIME_CLOCK", "RTC")
            .replace("SFR_SPECIAL_FUNCTION_REGISTERS", "SFR")
            .replace("PMM_POWER_MANAGEMENT_SYSTEM", "PMM")
            .replace("RC_RAM_CONTROL_MODULE", "RC")
            .replace("UCS_UNIFIED_SYSTEM_CLOCK", "UCS")
            .replace("SYS_SYSTEM_MODULE", "SYS")
            .replace("MPY_16_MULTIPLIER_16_BIT_MODE", "MPY_16")
            .replace("MPY_32_MULTIPLIER_32_BIT_MODE", "MPY_32")
            .replace("CS_CLOCK_SYSTEM", "CS")
    }
}

fn get_bits(offset: u32, width: u32) -> u64 {
    let mask = (1u64 << width) - 1;
    mask << offset
}

#[test]
fn test_get_bits() {
    let v = get_bits(0, 2);
    assert_eq!(v, 3);
}

pub fn build_svd_device(
    dev: &dslite_parser::Device,
    interrupts: &header_parser::Interrupts,
) -> svd::Device {
    let mut peripherals = Vec::new();
    for (_, m) in &dev.modules {
        let base_address = m.registers.iter().map(|r| r.offset).min().unwrap_or(0) & (!1);
        assert!(base_address < (1 << 16));
        assert!(base_address % 2 == 0);

        let mut registers = Vec::new();
        for reg in &m.registers {
            let offset = reg.offset - base_address;
            assert!(offset < (1 << 16));
            if reg.width > 1 {
                assert!(offset % 2 == 0);
            }

            let reset_mask = match reg.width {
                1 => 0xff,
                2 => 0xffff,
                4 => 0xffff_ffff,
                _ => panic!("invalid width"),
            };

            let mut fields = Vec::new();
            for f in &reg.fields {
                let mut enums = Vec::new();
                for e in &f.enums {
                    let ev = svd::EnumeratedValue {
                        name: e.name.clone(),
                        description: Some(e.description.clone()),
                        value: Some(e.value),
                        is_default: None,
                    };
                    enums.push(ev);
                }

                let field_constraint = None;
                if enums.is_empty() && f.width > 1 {
                    eprintln!("warning: no enums for field {}", f.name);
                    // field_constraint =
                    //     Some(svd::WriteConstraint::Range(svd::WriteConstraintRange {
                    //         min: 0,
                    //         max: (1 << f.width) - 1,
                    //     }));
                }
                let access = match f.rwa {
                    dslite_parser::Access::ReadWrite => svd::Access::ReadWrite,
                    dslite_parser::Access::ReadOnly => svd::Access::ReadOnly,
                    dslite_parser::Access::WriteOnly => svd::Access::WriteOnly,
                };

                let field = svd::Field {
                    name: f.name.clone(),
                    description: Some(f.description.clone()),
                    bit_range: svd::BitRange {
                        offset: f.offset,
                        width: f.width,
                    },
                    access: Some(access),
                    enumerated_values: vec![svd::EnumeratedValues {
                        name: None,
                        usage: None,
                        derived_from: None,
                        values: enums,
                    }],
                    write_constraint: field_constraint,
                };
                fields.push(field);
            }

            let mut reg_constraint = None;
            if fields.is_empty() {
                eprintln!("warning: no fields in register {}", reg.name);
            }

            // check that fields have correct range
            let mut field_test = !get_bits(0, reg.width * 8);
            for f in &fields {
                let field_bits = get_bits(f.bit_range.offset, f.bit_range.width);
                if field_test & field_bits != 0 {
                    eprintln!("error: bad field {} in register {}", f.name, reg.name);
                }
                field_test |= field_bits;
            }

            // if all fields are single bit and cover the entire register
            // then allow to write any value
            if (fields.len() == reg.width as usize * 8)
                && (fields.iter().all(|f| f.bit_range.width == 1))
            {
                reg_constraint = Some(svd::WriteConstraint::Range(svd::WriteConstraintRange {
                    min: 0,
                    max: ((1u64 << (reg.width * 8)) - 1) as u32,
                }));
            }

            let ri = svd::RegisterInfo {
                name: reg.name.fix_name(),
                description: reg.description.clone(),
                address_offset: offset,
                size: Some(reg.width * 8),
                access: None,
                reset_value: None,
                reset_mask: Some(reset_mask),
                fields: if fields.is_empty() {
                    None
                } else {
                    Some(fields)
                },
                write_constraint: reg_constraint,
                alternate_register: reg.alternate.clone(),
            };
            registers.push(svd::Register::Single(ri));
        }

        let p = svd::Peripheral {
            name: m.name.fix_name(),
            group_name: None,
            description: Some(m.description.clone()),
            base_address: base_address,
            interrupt: Vec::new(),
            registers: Some(registers),
            derived_from: None,
        };
        peripherals.push(p);
    }

    if !interrupts.vectors.is_empty() {
        peripherals.push(svd::Peripheral {
            name: "_INTERRUPTS".to_owned(),
            group_name: None,
            description: None,
            base_address: interrupts.base_offset,
            interrupt: interrupts
                .vectors
                .iter()
                .map(|int| svd::Interrupt {
                    name: int.name.clone(),
                    description: Some(int.description.clone()),
                    value: int.value,
                })
                .collect(),
            registers: None,
            derived_from: None,
        });
    }

    svd::Device {
        name: dev.name.fix_name(),
        peripherals: peripherals,
        defaults: svd::Defaults {
            size: Some(2),
            reset_value: Some(0),
            reset_mask: None,
            access: Some(svd::Access::ReadWrite),
        },
    }
}
<commit_msg>Do not generate empty enumerated values which violate the spec.<commit_after>use dslite_parser;
use header_parser;
use inflector::Inflector;
use svd;

trait StringEx {
    fn fix_name(&self) -> String;
}

impl StringEx for String {
    fn fix_name(&self) -> String {
        self.to_screaming_snake_case()
            // Some fixups to make the names look nicer
            .replace("RTC_REAL_TIME_CLOCK", "RTC")
            .replace("SFR_SPECIAL_FUNCTION_REGISTERS", "SFR")
            .replace("PMM_POWER_MANAGEMENT_SYSTEM", "PMM")
            .replace("RC_RAM_CONTROL_MODULE", "RC")
            .replace("UCS_UNIFIED_SYSTEM_CLOCK", "UCS")
            .replace("SYS_SYSTEM_MODULE", "SYS")
            .replace("MPY_16_MULTIPLIER_16_BIT_MODE", "MPY_16")
            .replace("MPY_32_MULTIPLIER_32_BIT_MODE", "MPY_32")
            .replace("CS_CLOCK_SYSTEM", "CS")
    }
}

fn get_bits(offset: u32, width: u32) -> u64 {
    let mask = (1u64 << width) - 1;
    mask << offset
}

#[test]
fn test_get_bits() {
    let v = get_bits(0, 2);
    assert_eq!(v, 3);
}

pub fn build_svd_device(
    dev: &dslite_parser::Device,
    interrupts: &header_parser::Interrupts,
) -> svd::Device {
    let mut peripherals = Vec::new();
    for (_, m) in &dev.modules {
        let base_address = m.registers.iter().map(|r| r.offset).min().unwrap_or(0) & (!1);
        assert!(base_address < (1 << 16));
        assert!(base_address % 2 == 0);

        let mut registers = Vec::new();
        for reg in &m.registers {
            let offset = reg.offset - base_address;
            assert!(offset < (1 << 16));
            if reg.width > 1 {
                assert!(offset % 2 == 0);
            }

            let reset_mask = match reg.width {
                1 => 0xff,
                2 => 0xffff,
                4 => 0xffff_ffff,
                _ => panic!("invalid width"),
            };

            let mut fields = Vec::new();
            for f in &reg.fields {
                let mut enums = Vec::new();
                for e in &f.enums {
                    let ev = svd::EnumeratedValue {
                        name: e.name.clone(),
                        description: Some(e.description.clone()),
                        value: Some(e.value),
                        is_default: None,
                    };
                    enums.push(ev);
                }

                let field_constraint = None;
                if enums.is_empty() && f.width > 1 {
                    eprintln!("warning: no enums for field {}", f.name);
                    // field_constraint =
                    //     Some(svd::WriteConstraint::Range(svd::WriteConstraintRange {
                    //         min: 0,
                    //         max: (1 << f.width) - 1,
                    //     }));
                }
                let access = match f.rwa {
                    dslite_parser::Access::ReadWrite => svd::Access::ReadWrite,
                    dslite_parser::Access::ReadOnly => svd::Access::ReadOnly,
                    dslite_parser::Access::WriteOnly => svd::Access::WriteOnly,
                };

                let field = svd::Field {
                    name: f.name.clone(),
                    description: Some(f.description.clone()),
                    bit_range: svd::BitRange {
                        offset: f.offset,
                        width: f.width,
                    },
                    access: Some(access),
                    enumerated_values: if enums.len() != 0 {
                        vec![svd::EnumeratedValues {
                            name: None,
                            usage: None,
                            derived_from: None,
                            values: enums,
                        }]
                    } else {
                        vec![]
                    },
                    write_constraint: field_constraint,
                };
                fields.push(field);
            }

            let mut reg_constraint = None;
            if fields.is_empty() {
                eprintln!("warning: no fields in register {}", reg.name);
            }

            // check that fields have correct range
            let mut field_test = !get_bits(0, reg.width * 8);
            for f in &fields {
                let field_bits = get_bits(f.bit_range.offset, f.bit_range.width);
                if field_test & field_bits != 0 {
                    eprintln!("error: bad field {} in register {}", f.name, reg.name);
                }
                field_test |= field_bits;
            }

            // if all fields are single bit and cover the entire register
            // then allow to write any value
            if (fields.len() == reg.width as usize * 8)
                && (fields.iter().all(|f| f.bit_range.width == 1))
            {
                reg_constraint = Some(svd::WriteConstraint::Range(svd::WriteConstraintRange {
                    min: 0,
                    max: ((1u64 << (reg.width * 8)) - 1) as u32,
                }));
            }

            let ri = svd::RegisterInfo {
                name: reg.name.fix_name(),
                description: reg.description.clone(),
                address_offset: offset,
                size: Some(reg.width * 8),
                access: None,
                reset_value: None,
                reset_mask: Some(reset_mask),
                fields: if fields.is_empty() {
                    None
                } else {
                    Some(fields)
                },
                write_constraint: reg_constraint,
                alternate_register: reg.alternate.clone(),
            };
            registers.push(svd::Register::Single(ri));
        }

        let p = svd::Peripheral {
            name: m.name.fix_name(),
            group_name: None,
            description: Some(m.description.clone()),
            base_address: base_address,
            interrupt: Vec::new(),
            registers: Some(registers),
            derived_from: None,
        };
        peripherals.push(p);
    }

    if !interrupts.vectors.is_empty() {
        peripherals.push(svd::Peripheral {
            name: "_INTERRUPTS".to_owned(),
            group_name: None,
            description: None,
            base_address: interrupts.base_offset,
            interrupt: interrupts
                .vectors
                .iter()
                .map(|int| svd::Interrupt {
                    name: int.name.clone(),
                    description: Some(int.description.clone()),
                    value: int.value,
                })
                .collect(),
            registers: None,
            derived_from: None,
        });
    }

    svd::Device {
        name: dev.name.fix_name(),
        peripherals: peripherals,
        defaults: svd::Defaults {
            size: Some(2),
            reset_value: Some(0),
            reset_mask: None,
            access: Some(svd::Access::ReadWrite),
        },
    }
}
<|endoftext|>


<commit_before>//! Disassembly calling function for most targets.

use crate::Function;
use std::{collections::HashSet, env, process::Command, str};

// Extracts the "shim" name from the `symbol`.
fn normalize(mut symbol: &str) -> String {
    // Remove trailing colon:
    if symbol.ends_with(':') {
        symbol = &symbol[..symbol.len() - 1];
    }
    if symbol.ends_with('>') {
        symbol = &symbol[..symbol.len() - 1];
    }
    if let Some(idx) = symbol.find('<') {
        symbol = &symbol[idx + 1..];
    }

    let mut symbol = rustc_demangle::demangle(symbol).to_string();
    symbol = match symbol.rfind("::h") {
        Some(i) => symbol[..i].to_string(),
        None => symbol.to_string(),
    };

    // Remove Rust paths
    if let Some(last_colon) = symbol.rfind(':') {
        symbol = (&symbol[last_colon + 1..]).to_string();
    }

    // Normalize to no leading underscore to handle platforms that may
    // inject extra ones in symbol names.
    while symbol.starts_with('_') {
        symbol.remove(0);
    }
    // Windows/x86 has a suffix such as @@4.
    if let Some(idx) = symbol.find("@@") {
        symbol = (&symbol[..idx]).to_string();
    }
    symbol
}

pub(crate) fn disassemble_myself() -> HashSet<Function> {
    let me = env::current_exe().expect("failed to get current exe");

    let disassembly = if cfg!(target_os = "windows") && cfg!(target_env = "msvc") {
        let target = if cfg!(target_arch = "x86_64") {
            "x86_64-pc-windows-msvc"
        } else if cfg!(target_arch = "x86") {
            "i686-pc-windows-msvc"
        } else {
            panic!("disassembly unimplemented")
        };
        let mut cmd = cc::windows_registry::find(target, "dumpbin.exe")
            .expect("failed to find `dumpbin` tool");
        let output = cmd
            .arg("/DISASM")
            .arg(&me)
            .output()
            .expect("failed to execute dumpbin");
        println!(
            "{}\n{}",
            output.status,
            String::from_utf8_lossy(&output.stderr)
        );
        assert!(output.status.success());
        // Windows does not return valid UTF-8 output:
        String::from_utf8_lossy(Vec::leak(output.stdout))
    } else if cfg!(target_os = "windows") {
        panic!("disassembly unimplemented")
    } else {
        let objdump = env::var("OBJDUMP").unwrap_or_else(|_| "objdump".to_string());
        let add_args = if cfg!(target_os = "macos") && cfg!(target_arch = "aarch64") {
            // Target features need to be enabled for LLVM objdump on Macos ARM64
            vec!["--mattr=+crc,+crypto,+tme"]
        } else {
            vec![]
        };
        let output = Command::new(objdump.clone())
            .arg("--disassemble")
            .arg("--no-show-raw-insn")
            .args(add_args)
            .arg(&me)
            .output()
            .unwrap_or_else(|_| panic!("failed to execute objdump. OBJDUMP={}", objdump));
        println!(
            "{}\n{}",
            output.status,
            String::from_utf8_lossy(&output.stderr)
        );
        assert!(output.status.success());

        String::from_utf8_lossy(Vec::leak(output.stdout))
    };

    parse(&disassembly)
}

fn parse(output: &str) -> HashSet<Function> {
    let mut lines = output.lines();

    println!(
        "First 100 lines of the disassembly input containing {} lines:",
        lines.clone().count()
    );
    for line in output.lines().take(100) {
        println!("{}", line);
    }

    let mut functions = HashSet::new();
    let mut cached_header = None;
    while let Some(header) = cached_header.take().or_else(|| lines.next()) {
        if !header.ends_with(':') || !header.contains("stdarch_test_shim") {
            continue;
        }
        eprintln!("header: {}", header);
        let symbol = normalize(header);
        eprintln!("normalized symbol: {}", symbol);
        let mut instructions = Vec::new();
        while let Some(instruction) = lines.next() {
            if instruction.ends_with(':') {
                cached_header = Some(instruction);
                break;
            }
            if instruction.is_empty() {
                cached_header = None;
                break;
            }
            let parts = if cfg!(target_os = "macos") {
                // Each line of instructions should look like:
                //
                //      $addr    $instruction...
                instruction
                    .split_whitespace()
                    .skip(1)
                    .map(std::string::ToString::to_string)
                    .collect::<Vec<String>>()
            } else if cfg!(target_env = "msvc") {
                // Each line looks like:
                //
                // >  $addr: ab cd ef     $instr..
                // >         00 12          # this line os optional
                if instruction.starts_with("       ") {
                    continue;
                }
                instruction
                    .split_whitespace()
                    .skip(1)
                    .skip_while(|s| s.len() == 2 && usize::from_str_radix(s, 16).is_ok())
                    .map(std::string::ToString::to_string)
                    .skip_while(|s| *s == "lock") // skip x86-specific prefix
                    .collect::<Vec<String>>()
            } else {
                // objdump with --no-show-raw-insn
                // Each line of instructions should look like:
                //
                //      $rel_offset:       $instruction...
                instruction
                    .split_whitespace()
                    .skip(1)
                    .skip_while(|s| *s == "lock") // skip x86-specific prefix
                    .map(std::string::ToString::to_string)
                    .collect::<Vec<String>>()
            };
            instructions.push(parts.join(" "));
        }
        let function = Function {
            name: symbol,
            instrs: instructions,
        };
        assert!(functions.insert(function));
    }

    eprintln!("all found functions dump:");
    for k in &functions {
        eprintln!("  f: {}", k.name);
    }

    functions
}
<commit_msg>using v8.6a target feature to cover more instructions<commit_after>//! Disassembly calling function for most targets.

use crate::Function;
use std::{collections::HashSet, env, process::Command, str};

// Extracts the "shim" name from the `symbol`.
fn normalize(mut symbol: &str) -> String {
    // Remove trailing colon:
    if symbol.ends_with(':') {
        symbol = &symbol[..symbol.len() - 1];
    }
    if symbol.ends_with('>') {
        symbol = &symbol[..symbol.len() - 1];
    }
    if let Some(idx) = symbol.find('<') {
        symbol = &symbol[idx + 1..];
    }

    let mut symbol = rustc_demangle::demangle(symbol).to_string();
    symbol = match symbol.rfind("::h") {
        Some(i) => symbol[..i].to_string(),
        None => symbol.to_string(),
    };

    // Remove Rust paths
    if let Some(last_colon) = symbol.rfind(':') {
        symbol = (&symbol[last_colon + 1..]).to_string();
    }

    // Normalize to no leading underscore to handle platforms that may
    // inject extra ones in symbol names.
    while symbol.starts_with('_') {
        symbol.remove(0);
    }
    // Windows/x86 has a suffix such as @@4.
    if let Some(idx) = symbol.find("@@") {
        symbol = (&symbol[..idx]).to_string();
    }
    symbol
}

pub(crate) fn disassemble_myself() -> HashSet<Function> {
    let me = env::current_exe().expect("failed to get current exe");

    let disassembly = if cfg!(target_os = "windows") && cfg!(target_env = "msvc") {
        let target = if cfg!(target_arch = "x86_64") {
            "x86_64-pc-windows-msvc"
        } else if cfg!(target_arch = "x86") {
            "i686-pc-windows-msvc"
        } else {
            panic!("disassembly unimplemented")
        };
        let mut cmd = cc::windows_registry::find(target, "dumpbin.exe")
            .expect("failed to find `dumpbin` tool");
        let output = cmd
            .arg("/DISASM")
            .arg(&me)
            .output()
            .expect("failed to execute dumpbin");
        println!(
            "{}\n{}",
            output.status,
            String::from_utf8_lossy(&output.stderr)
        );
        assert!(output.status.success());
        // Windows does not return valid UTF-8 output:
        String::from_utf8_lossy(Vec::leak(output.stdout))
    } else if cfg!(target_os = "windows") {
        panic!("disassembly unimplemented")
    } else {
        let objdump = env::var("OBJDUMP").unwrap_or_else(|_| "objdump".to_string());
        let add_args = if cfg!(target_os = "macos") && cfg!(target_arch = "aarch64") {
            // Target features need to be enabled for LLVM objdump on Macos ARM64
            vec!["--mattr=+v8.6a,+crypto,+tme"]
        } else {
            vec![]
        };
        let output = Command::new(objdump.clone())
            .arg("--disassemble")
            .arg("--no-show-raw-insn")
            .args(add_args)
            .arg(&me)
            .output()
            .unwrap_or_else(|_| panic!("failed to execute objdump. OBJDUMP={}", objdump));
        println!(
            "{}\n{}",
            output.status,
            String::from_utf8_lossy(&output.stderr)
        );
        assert!(output.status.success());

        String::from_utf8_lossy(Vec::leak(output.stdout))
    };

    parse(&disassembly)
}

fn parse(output: &str) -> HashSet<Function> {
    let mut lines = output.lines();

    println!(
        "First 100 lines of the disassembly input containing {} lines:",
        lines.clone().count()
    );
    for line in output.lines().take(100) {
        println!("{}", line);
    }

    let mut functions = HashSet::new();
    let mut cached_header = None;
    while let Some(header) = cached_header.take().or_else(|| lines.next()) {
        if !header.ends_with(':') || !header.contains("stdarch_test_shim") {
            continue;
        }
        eprintln!("header: {}", header);
        let symbol = normalize(header);
        eprintln!("normalized symbol: {}", symbol);
        let mut instructions = Vec::new();
        while let Some(instruction) = lines.next() {
            if instruction.ends_with(':') {
                cached_header = Some(instruction);
                break;
            }
            if instruction.is_empty() {
                cached_header = None;
                break;
            }
            let parts = if cfg!(target_os = "macos") {
                // Each line of instructions should look like:
                //
                //      $addr    $instruction...
                instruction
                    .split_whitespace()
                    .skip(1)
                    .map(std::string::ToString::to_string)
                    .collect::<Vec<String>>()
            } else if cfg!(target_env = "msvc") {
                // Each line looks like:
                //
                // >  $addr: ab cd ef     $instr..
                // >         00 12          # this line os optional
                if instruction.starts_with("       ") {
                    continue;
                }
                instruction
                    .split_whitespace()
                    .skip(1)
                    .skip_while(|s| s.len() == 2 && usize::from_str_radix(s, 16).is_ok())
                    .map(std::string::ToString::to_string)
                    .skip_while(|s| *s == "lock") // skip x86-specific prefix
                    .collect::<Vec<String>>()
            } else {
                // objdump with --no-show-raw-insn
                // Each line of instructions should look like:
                //
                //      $rel_offset:       $instruction...
                instruction
                    .split_whitespace()
                    .skip(1)
                    .skip_while(|s| *s == "lock") // skip x86-specific prefix
                    .map(std::string::ToString::to_string)
                    .collect::<Vec<String>>()
            };
            instructions.push(parts.join(" "));
        }
        let function = Function {
            name: symbol,
            instrs: instructions,
        };
        assert!(functions.insert(function));
    }

    eprintln!("all found functions dump:");
    for k in &functions {
        eprintln!("  f: {}", k.name);
    }

    functions
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The pnet_macros crate provides the `#[packet]` macro and compiler plugin, which is used to
//! specify the format of on-the-wire packets, and automatically generate zero-copy accessors and
//! mutators for the fields. It is used as follows:
//!
//! ```
//! #![feature(core, collections, custom_attribute, plugin)]
//! #![plugin(pnet_macros)]
//!
//! extern crate pnet;
//!
//! /// This module contains a list of type aliases which may be used
//! use pnet_macros::types::{u4, u12be};
//!
//! /// Packets are specified in the same way as normal Rust structs, but with a `#[packet]`
//! /// attribute.
//! #[packet]
//! pub struct Example {
//!     /// This is a simple field which contains a 4-bit, unsigned integer.
//!     /// Note that `u4` is simply an alias for `u8` - the name is a hint
//!     /// to the compiler plugin, it is NOT a usable 4 bit type!
//!     simple_field1: u4,
//!
//!     /// This specifies that `simple_field2` should be a 12-bit field,
//!     /// with bits stored in big endian
//!     simple_field2: u12be,
//!
//!     /// All packets must specify a `#[payload]`, which should be a
//!     /// `Vec<u8>`. This represents the packet's payload, for example in
//!     /// an IPv4 packet, the payload could be a UDP packet, or in a UDP
//!     /// packet the payload could be the application data. All the
//!     /// remaining space in the packet is considered to be the payload
//!     /// (this doesn't have to be the case, see the documentation for
//!     /// `#[payload]` below.
//!     #[payload]
//!     payload: Vec<u8>
//! }
//! ```
//! A number of things will then be generated. You can see this in action in the documentation and
//! source of each of the packet types in the `pnet::packet` module. Things generated include
//! (assuming the `Example` struct from above):
//!
//!  * An `ExamplePacket<'p>` structure, which is used for receiving packets on the network.
//!    This structure contains:
//!      - A method, `pub fn new<'p>(packet: &'p [u8]) -> ExamplePacket<'p>`, used for the
//!        construction of an `ExamplePacket`, given a buffer to store it. The buffer should be
//!        long enough to contain all the fields in the packet.
//!      - A method, `pub fn to_immutable<'p>(&'p self) -> ExamplePacket<'p>`, which is simply an
//!        identity function. It exists for consistency with `MutableExamplePacket`.
//!      - A number of accessor methods, of the form `pub get_{field_name}(&self) -> {field_type}`,
//!        which will retreive the host representation of the on-the-wire value.
//!  * A `MutableExamplePacket<'p>` structure, which is used when sending packets on the network.
//!    This structure contains:
//!      - A method, `pub fn new<'p>(packet: &'p mut [u8]) -> MutableExamplePacket<'p>`, used for
//!        the construction of a `MutableExamplePacket`, given a buffer to store it. The buffer
//!        should be long enough to contain all the fields in the packet.
//!      - A method, `pub fn to_immutable<'p>(&'p self) -> ExamplePacket<'p>`, which converts from
//!        a `MutableExamplePacket` to an `ExamplePacket`
//!      - A method, `pub fn populate(&mut self, packet: Example)`, which, given an `Example`
//!        struct, will populate the `MutableExamplePacket` with the values from the `Example`
//!        struct.
//!      - A number of accessor methods, of the form `pub get_{field_name}(&self) -> {field_type}`,
//!        which will retreive the host representation of the on-the-wire value.
//!      - A number of mutator methods, of the form `pub set_{field_name}(&mut self,
//!        val: {field_type})`, which will take a host value, convert it to the required
//!        on-the-wire format, and store it in the buffer which backs the `MutableExamplePacket`.
//!  * A number of trait implementations for each of the `MutableExamplePacket` and `ExamplePacket`
//!    structures. These include:
//!      - `pnet::packet::Packet` (`ExamplePacket` and `MutableExamplePacket`)
//!      - `pnet::packet::MutablePacket` (`MutableExamplePacket` only)
//!      - `std::fmt::Debug` (`ExamplePacket` and `MutableExamplePacket`)
//!      - `pnet::packet::FromPacket` (`ExamplePacket` and `MutableExamplePacket`)
//!      - `pnet::packet::PacketSize` (`ExamplePacket` and `MutableExamplePacket`)
//!  * An `ExampleIterator` structure, which implements `std::iter::Iterator`, to allow iterating
//!    over vectors of `ExamplePacket` contained within another packet. Used internally.
//!
//! ## Attributes
//!
//! There are a number of attributes which fields may have, these include:
//!
//!  * \#[length_fn = "function_name"]
//!
//!    This attribute is used to enable variable length fields. To specify a variable length field,
//!    it should have the type `Vec<T>`. It must have the `#[length_fn]` attribute, which specifies
//!    a function name to calculate the length of the field. The signature for the length function
//!    should be `fn {function_name}<'a>(example_packet: &ExamplePacket<'a>) -> usize`,
//!    substituting `&ExamplePacket<'a>` for the appropriately named packet type for your
//!    structure. You may access whichever fields are required to calculate the length of the
//!    field. The returned value should be a number of bytes that the field uses.
//!
//!    The type contained in the vector may either be one of the primitive types specified in
//!    `pnet_macros::types`, or another structure marked with #[packet], for example
//!    `Vec<Example>`.
//!
//!  * \#[payload]
//!
//!    This attribute specifies the payload associated with the packet. This should specify the
//!    data associated with the packet. It may be used in two places:
//!      - The last field in the packet, in which case it is assumed to use the remaining length of
//!        the buffer containing the packet
//!      - Another location in the packet, in which case the `#[length_fn]` attribute must also be
//!        specified to give the length of the payload.
//!    If the packet has no payload, you must still specify this attribute, but you can provide a
//!    `#[length_fn]` attribute returning zero.
//!
//!  * \#[construct_with(<primitive type>, ...)]
//!
//!    Unfortunately, compiler plugins do not currently have access to type information during the
//!    decoration stage (where all of the above is generated), so this attribute is required. This
//!    must be used for all fields which are neither primitive types, nor vectors of primitive
//!    types. Three things are required when using `#[construct_with]`:
//!      - The field type must have a method `new`, which takes one or more parameters of primitive
//!        types.
//!      - The field must be annotated with the `#[construct_with(...)]` attribute, specifying a
//!        list of types identical to those taken by the `new` method.
//!      - The `pnet::packet::ToPrimitiveValues` trait must be implemented for the field type,
//!        which must return a tuple of the primitive types specified in the parameters to the
//!        `#[construct_with(...)]` attribute, and in the `new` method.

#![warn(missing_docs)]

#![feature(core, collections, plugin_registrar, quote, rustc_private)]

extern crate syntax;
extern crate regex;
#[macro_use] extern crate rustc;

use rustc::lint::{LintPassObject};
use rustc::plugin::Registry;

use syntax::ast;
use syntax::codemap::{Span};
use syntax::parse::token;
use syntax::ext::base::{Decorator, ExtCtxt, Modifier};
use syntax::ptr::P;

mod decorator;
mod lint;
mod util;

pub mod types;

/// Replace the #[packet] attribute with internal attributes
///
/// The #[packet] attribute is consumed, so we replace it with two internal attributes,
/// #[_packet_generator], which is used to generate the packet implementations, and
/// #[_packet_lint], which is used to trigger linting.
fn packet_modifier(ecx: &mut ExtCtxt,
                   _span: Span,
                   _meta_item: &ast::MetaItem,
                   item: P<ast::Item>) -> P<ast::Item> {
    let mut new_item = (*item).clone();

    new_item.attrs.push(quote_attr!(ecx, #[_packet_lint]));
    new_item.attrs.push(quote_attr!(ecx, #[_packet_generator]));
    new_item.attrs.push(quote_attr!(ecx, #[derive(Clone, Debug)]));
    new_item.attrs.push(quote_attr!(ecx, #[allow(unused_attributes)]));

    P(new_item)
}

/// The entry point for the syntax extension
///
/// This registers each part of the plugin with the compiler. There are three parts, a modifier, to
/// add additional attributes to the original structure, the decorator, to generate the various
/// required structures and method, and a lint pass, which performs additional validation which
/// requires type information.
#[plugin_registrar]
pub fn plugin_registrar(registry: &mut Registry) {
    registry.register_syntax_extension(token::intern("packet"),
                                       Modifier(Box::new(packet_modifier)));
    registry.register_syntax_extension(token::intern("_packet_generator"),
                                       Decorator(Box::new(decorator::generate_packet)));

    registry.register_lint_pass(Box::new(lint::PacketPass) as LintPassObject);
}
<commit_msg>Add `#[length]` to documentation<commit_after>// Copyright (c) 2015 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The pnet_macros crate provides the `#[packet]` macro and compiler plugin, which is used to
//! specify the format of on-the-wire packets, and automatically generate zero-copy accessors and
//! mutators for the fields. It is used as follows:
//!
//! ```
//! #![feature(core, collections, custom_attribute, plugin)]
//! #![plugin(pnet_macros)]
//!
//! extern crate pnet;
//!
//! /// This module contains a list of type aliases which may be used
//! use pnet_macros::types::{u4, u12be};
//!
//! /// Packets are specified in the same way as normal Rust structs, but with a `#[packet]`
//! /// attribute.
//! #[packet]
//! pub struct Example {
//!     /// This is a simple field which contains a 4-bit, unsigned integer.
//!     /// Note that `u4` is simply an alias for `u8` - the name is a hint
//!     /// to the compiler plugin, it is NOT a usable 4 bit type!
//!     simple_field1: u4,
//!
//!     /// This specifies that `simple_field2` should be a 12-bit field,
//!     /// with bits stored in big endian
//!     simple_field2: u12be,
//!
//!     /// All packets must specify a `#[payload]`, which should be a
//!     /// `Vec<u8>`. This represents the packet's payload, for example in
//!     /// an IPv4 packet, the payload could be a UDP packet, or in a UDP
//!     /// packet the payload could be the application data. All the
//!     /// remaining space in the packet is considered to be the payload
//!     /// (this doesn't have to be the case, see the documentation for
//!     /// `#[payload]` below.
//!     #[payload]
//!     payload: Vec<u8>
//! }
//! ```
//! A number of things will then be generated. You can see this in action in the documentation and
//! source of each of the packet types in the `pnet::packet` module. Things generated include
//! (assuming the `Example` struct from above):
//!
//!  * An `ExamplePacket<'p>` structure, which is used for receiving packets on the network.
//!    This structure contains:
//!      - A method, `pub fn new<'p>(packet: &'p [u8]) -> ExamplePacket<'p>`, used for the
//!        construction of an `ExamplePacket`, given a buffer to store it. The buffer should be
//!        long enough to contain all the fields in the packet.
//!      - A method, `pub fn to_immutable<'p>(&'p self) -> ExamplePacket<'p>`, which is simply an
//!        identity function. It exists for consistency with `MutableExamplePacket`.
//!      - A number of accessor methods, of the form `pub get_{field_name}(&self) -> {field_type}`,
//!        which will retreive the host representation of the on-the-wire value.
//!  * A `MutableExamplePacket<'p>` structure, which is used when sending packets on the network.
//!    This structure contains:
//!      - A method, `pub fn new<'p>(packet: &'p mut [u8]) -> MutableExamplePacket<'p>`, used for
//!        the construction of a `MutableExamplePacket`, given a buffer to store it. The buffer
//!        should be long enough to contain all the fields in the packet.
//!      - A method, `pub fn to_immutable<'p>(&'p self) -> ExamplePacket<'p>`, which converts from
//!        a `MutableExamplePacket` to an `ExamplePacket`
//!      - A method, `pub fn populate(&mut self, packet: Example)`, which, given an `Example`
//!        struct, will populate the `MutableExamplePacket` with the values from the `Example`
//!        struct.
//!      - A number of accessor methods, of the form `pub get_{field_name}(&self) -> {field_type}`,
//!        which will retreive the host representation of the on-the-wire value.
//!      - A number of mutator methods, of the form `pub set_{field_name}(&mut self,
//!        val: {field_type})`, which will take a host value, convert it to the required
//!        on-the-wire format, and store it in the buffer which backs the `MutableExamplePacket`.
//!  * A number of trait implementations for each of the `MutableExamplePacket` and `ExamplePacket`
//!    structures. These include:
//!      - `pnet::packet::Packet` (`ExamplePacket` and `MutableExamplePacket`)
//!      - `pnet::packet::MutablePacket` (`MutableExamplePacket` only)
//!      - `std::fmt::Debug` (`ExamplePacket` and `MutableExamplePacket`)
//!      - `pnet::packet::FromPacket` (`ExamplePacket` and `MutableExamplePacket`)
//!      - `pnet::packet::PacketSize` (`ExamplePacket` and `MutableExamplePacket`)
//!  * An `ExampleIterator` structure, which implements `std::iter::Iterator`, to allow iterating
//!    over vectors of `ExamplePacket` contained within another packet. Used internally.
//!
//! ## Attributes
//!
//! There are a number of attributes which fields may have, these include:
//!
//!  * \#[length_fn = "function_name"]
//!
//!    This attribute is used to enable variable length fields. To specify a variable length field,
//!    it should have the type `Vec<T>`. It must have the `#[length_fn]` (or #[length]) attribute,
//!    which specifies a function name to calculate the length of the field. The signature for the
//!    length function should be
//!    `fn {function_name}<'a>(example_packet: &ExamplePacket<'a>) -> usize`, substituting
//!    `&ExamplePacket<'a>` for the appropriately named packet type for your structure. You may
//!    access whichever fields are required to calculate the length of the field. The returned
//!    value should be a number of bytes that the field uses.
//!
//!    The type contained in the vector may either be one of the primitive types specified in
//!    `pnet_macros::types`, or another structure marked with #[packet], for example
//!    `Vec<Example>`.
//!
//!  * \#[length = "arithmetic expression"]
//!
//!    This attribute is used to enable variable length fields. To specify a variable length field,
//!    it should have the type `Vec<T>`. It must have the `#[length]` (or #[length_fn]) attribute,
//!    which specifies an arithmetic expression to calculate the length of the field. Only field
//!    names, constants, integers, basic arithmetic expressions (+ - * / %) and parentheses are
//!    in the expression. An example would be `#[length = "field_name + CONSTANT - 4]`.
//!
//!    The type contained in the vector may either be one of the primitive types specified in
//!    `pnet_macros::types`, or another structure marked with #[packet], for example
//!    `Vec<Example>`.
//!
//!  * \#[payload]
//!
//!    This attribute specifies the payload associated with the packet. This should specify the
//!    data associated with the packet. It may be used in two places:
//!      - The last field in the packet, in which case it is assumed to use the remaining length of
//!        the buffer containing the packet
//!      - Another location in the packet, in which case the `#[length_fn]` attribute must also be
//!        specified to give the length of the payload.
//!    If the packet has no payload, you must still specify this attribute, but you can provide a
//!    `#[length_fn]` attribute returning zero.
//!
//!  * \#[construct_with(<primitive type>, ...)]
//!
//!    Unfortunately, compiler plugins do not currently have access to type information during the
//!    decoration stage (where all of the above is generated), so this attribute is required. This
//!    must be used for all fields which are neither primitive types, nor vectors of primitive
//!    types. Three things are required when using `#[construct_with]`:
//!      - The field type must have a method `new`, which takes one or more parameters of primitive
//!        types.
//!      - The field must be annotated with the `#[construct_with(...)]` attribute, specifying a
//!        list of types identical to those taken by the `new` method.
//!      - The `pnet::packet::ToPrimitiveValues` trait must be implemented for the field type,
//!        which must return a tuple of the primitive types specified in the parameters to the
//!        `#[construct_with(...)]` attribute, and in the `new` method.

#![warn(missing_docs)]

#![feature(core, collections, plugin_registrar, quote, rustc_private)]

extern crate syntax;
extern crate regex;
#[macro_use] extern crate rustc;

use rustc::lint::{LintPassObject};
use rustc::plugin::Registry;

use syntax::ast;
use syntax::codemap::{Span};
use syntax::parse::token;
use syntax::ext::base::{Decorator, ExtCtxt, Modifier};
use syntax::ptr::P;

mod decorator;
mod lint;
mod util;

pub mod types;

/// Replace the #[packet] attribute with internal attributes
///
/// The #[packet] attribute is consumed, so we replace it with two internal attributes,
/// #[_packet_generator], which is used to generate the packet implementations, and
/// #[_packet_lint], which is used to trigger linting.
fn packet_modifier(ecx: &mut ExtCtxt,
                   _span: Span,
                   _meta_item: &ast::MetaItem,
                   item: P<ast::Item>) -> P<ast::Item> {
    let mut new_item = (*item).clone();

    new_item.attrs.push(quote_attr!(ecx, #[_packet_lint]));
    new_item.attrs.push(quote_attr!(ecx, #[_packet_generator]));
    new_item.attrs.push(quote_attr!(ecx, #[derive(Clone, Debug)]));
    new_item.attrs.push(quote_attr!(ecx, #[allow(unused_attributes)]));

    P(new_item)
}

/// The entry point for the syntax extension
///
/// This registers each part of the plugin with the compiler. There are three parts, a modifier, to
/// add additional attributes to the original structure, the decorator, to generate the various
/// required structures and method, and a lint pass, which performs additional validation which
/// requires type information.
#[plugin_registrar]
pub fn plugin_registrar(registry: &mut Registry) {
    registry.register_syntax_extension(token::intern("packet"),
                                       Modifier(Box::new(packet_modifier)));
    registry.register_syntax_extension(token::intern("_packet_generator"),
                                       Decorator(Box::new(decorator::generate_packet)));

    registry.register_lint_pass(Box::new(lint::PacketPass) as LintPassObject);
}
<|endoftext|>


<commit_before>extern crate rusqlite;

use common::Link;
use common::types::StructType;
use common::types::FnType;
use common::Criteria;

use common::utils::load_file;

use super::base::Dao;
use super::base::SqlConfig;
use super::link_dao::LinkDao;

use self::rusqlite::Connection;

fn get_sql_config() -> SqlConfig {
    let db = match Connection::open_in_memory() {
        Ok(c) => c,
        Err(err) => panic!("OPEN TEST DB FAILED: {}", err),
    };

    let read_sql = String::from("res/sql/entity/read.sql");
    let delete_sql = String::from("res/sql/entity/delete.sql");
    let insert_sql = String::from("res/sql/entity/insert.sql");
    let update_sql = String::from("res/sql/entity/update.sql");
    let list_sql = String::from("res/sql/entity/list.sql");

    SqlConfig {
        connection: db,
        read_sql: load_file(&read_sql).unwrap(),
        delete_sql: load_file(&delete_sql).unwrap(),
        insert_sql: load_file(&insert_sql).unwrap(),
        update_sql: load_file(&update_sql).unwrap(),
        list_sql: load_file(&list_sql).unwrap(),
    }
}

fn init_test_db() -> SqlConfig {
    let config = get_sql_config();

    let init_sql = String::from("res/sql/entity/init.sql");
    let init_test_sql = String::from("res/sql/entity/init_test.sql");

    let init_query = load_file(&init_sql).unwrap();
    let init_test_query = load_file(&init_test_sql).unwrap();

    match config.connection.execute(init_query.as_str(), &[]) {
        Ok(_) => {}
        Err(err) => panic!("CREATE TEST TABLE FAILED: {}", err),
    };

    match config.connection.execute(init_test_query.as_str(), &[]) {
        Ok(_) => {}
        Err(err) => panic!("INIT TEST DB FAILED: {}", err),
    };

    config
}

#[test]
fn read() {
    let dao = LinkDao { config: init_test_db() };
    let id_test = 1;
    let link = Link {
        id: id_test,
        parent_id: String::from("1"),
        name: String::from("test_read"),
        url: String::from("test_url"),
        struct_type: StructType::Link,
        fn_type: FnType::None,
        rev_no: 0,
    };

    let read_c = Criteria::new();
    match dao.read(&read_c.id(id_test)) {
        Ok(l) => assert!(l == link),
        Err(err) => panic!("READ FAILED: {}", err),
    }
}

#[test]
fn insert() {
    let dao = LinkDao { config: init_test_db() };
    let id_test = 8;
    let link = Link {
        id: id_test,
        parent_id: String::from("3"),
        name: String::from("inserted"),
        url: String::from("test_url"),
        struct_type: StructType::Link,
        fn_type: FnType::None,
        rev_no: 0,
    };

    dao.insert(&link).unwrap();
    let read_c = Criteria::new();
    match dao.read(&read_c.id(id_test)) {
        Ok(l) => assert!(l == link),
        Err(err) => panic!("INSERT FAILED: {}", err),
    }
}

#[test]
#[should_panic]
fn test_delete() {
    let dao = LinkDao { config: init_test_db() };
    let id_test = 3;
    let crit = Criteria::new().id(id_test);

    dao.delete(&crit).unwrap();
    match dao.read(&crit) {
        Ok(_) => println!("{}", "IT NEVER PRINTS THIS"),
        Err(err) => panic!("{}", err),
    }
}

#[test]
fn test_update() {
    let dao = LinkDao { config: init_test_db() };
    let id_test = 4;

    let link = Link {
        id: id_test,
        parent_id: String::from("21"),
        name: String::from("updated"),
        url: String::from("test_url"),
        struct_type: StructType::Link,
        fn_type: FnType::None,
        rev_no: 0,
    };

    dao.update(link.clone()).unwrap();

    let read_c = Criteria::new();
    match dao.read(&read_c.id(id_test)) {
        Ok(l) => assert!(link == l),
        Err(err) => panic!("UPDATE FAILED: {}", err),
    }
}

#[test]
fn test_list() {
    let dao = LinkDao { config: init_test_db() };
    let inserted_name = String::from("inseted");

    let link1 = Link {
        id: -1,
        parent_id: String::from("3"),
        name: inserted_name.clone(),
        url: String::from("test_url1"),
        struct_type: StructType::Link,
        fn_type: FnType::None,
        rev_no: 0,
    };
    let link2 = Link {
        id: -1,
        parent_id: String::from("4"),
        name: inserted_name.clone(),
        url: String::from("test_url2"),
        struct_type: StructType::Node,
        fn_type: FnType::None,
        rev_no: 1,
    };

    dao.insert(&link1).unwrap();
    dao.insert(&link2).unwrap();

    let list_c = Criteria::new();
    match dao.list(&list_c.name(inserted_name)) {
        Ok(v) => assert!(v.len() == 2),
        Err(err) => panic!("LIST FAILED: {}", err),
    }
}
<commit_msg>fix string for list test<commit_after>extern crate rusqlite;

use common::Link;
use common::types::StructType;
use common::types::FnType;
use common::Criteria;

use common::utils::load_file;

use super::base::Dao;
use super::base::SqlConfig;
use super::link_dao::LinkDao;

use self::rusqlite::Connection;

fn get_sql_config() -> SqlConfig {
    let db = match Connection::open_in_memory() {
        Ok(c) => c,
        Err(err) => panic!("OPEN TEST DB FAILED: {}", err),
    };

    let read_sql = String::from("res/sql/entity/read.sql");
    let delete_sql = String::from("res/sql/entity/delete.sql");
    let insert_sql = String::from("res/sql/entity/insert.sql");
    let update_sql = String::from("res/sql/entity/update.sql");
    let list_sql = String::from("res/sql/entity/list.sql");

    SqlConfig {
        connection: db,
        read_sql: load_file(&read_sql).unwrap(),
        delete_sql: load_file(&delete_sql).unwrap(),
        insert_sql: load_file(&insert_sql).unwrap(),
        update_sql: load_file(&update_sql).unwrap(),
        list_sql: load_file(&list_sql).unwrap(),
    }
}

fn init_test_db() -> SqlConfig {
    let config = get_sql_config();

    let init_sql = String::from("res/sql/entity/init.sql");
    let init_test_sql = String::from("res/sql/entity/init_test.sql");

    let init_query = load_file(&init_sql).unwrap();
    let init_test_query = load_file(&init_test_sql).unwrap();

    match config.connection.execute(init_query.as_str(), &[]) {
        Ok(_) => {}
        Err(err) => panic!("CREATE TEST TABLE FAILED: {}", err),
    };

    match config.connection.execute(init_test_query.as_str(), &[]) {
        Ok(_) => {}
        Err(err) => panic!("INIT TEST DB FAILED: {}", err),
    };

    config
}

#[test]
fn read() {
    let dao = LinkDao { config: init_test_db() };
    let id_test = 1;
    let link = Link {
        id: id_test,
        parent_id: String::from("1"),
        name: String::from("test_read"),
        url: String::from("test_url"),
        struct_type: StructType::Link,
        fn_type: FnType::None,
        rev_no: 0,
    };

    let read_c = Criteria::new();
    match dao.read(&read_c.id(id_test)) {
        Ok(l) => assert!(l == link),
        Err(err) => panic!("READ FAILED: {}", err),
    }
}

#[test]
fn insert() {
    let dao = LinkDao { config: init_test_db() };
    let id_test = 8;
    let link = Link {
        id: id_test,
        parent_id: String::from("3"),
        name: String::from("inserted"),
        url: String::from("test_url"),
        struct_type: StructType::Link,
        fn_type: FnType::None,
        rev_no: 0,
    };

    dao.insert(&link).unwrap();
    let read_c = Criteria::new();
    match dao.read(&read_c.id(id_test)) {
        Ok(l) => assert!(l == link),
        Err(err) => panic!("INSERT FAILED: {}", err),
    }
}

#[test]
#[should_panic]
fn test_delete() {
    let dao = LinkDao { config: init_test_db() };
    let id_test = 3;
    let crit = Criteria::new().id(id_test);

    dao.delete(&crit).unwrap();
    match dao.read(&crit) {
        Ok(_) => println!("{}", "IT NEVER PRINTS THIS"),
        Err(err) => panic!("{}", err),
    }
}

#[test]
fn test_update() {
    let dao = LinkDao { config: init_test_db() };
    let id_test = 4;

    let link = Link {
        id: id_test,
        parent_id: String::from("21"),
        name: String::from("updated"),
        url: String::from("test_url"),
        struct_type: StructType::Link,
        fn_type: FnType::None,
        rev_no: 0,
    };

    dao.update(link.clone()).unwrap();

    let read_c = Criteria::new();
    match dao.read(&read_c.id(id_test)) {
        Ok(l) => assert!(link == l),
        Err(err) => panic!("UPDATE FAILED: {}", err),
    }
}

#[test]
fn test_list() {
    let dao = LinkDao { config: init_test_db() };
    let inserted_name = String::from("inserted");

    let link1 = Link {
        id: -1,
        parent_id: String::from("3"),
        name: inserted_name.clone(),
        url: String::from("test_url1"),
        struct_type: StructType::Link,
        fn_type: FnType::None,
        rev_no: 0,
    };
    let link2 = Link {
        id: -1,
        parent_id: String::from("4"),
        name: inserted_name.clone(),
        url: String::from("test_url2"),
        struct_type: StructType::Node,
        fn_type: FnType::None,
        rev_no: 1,
    };

    dao.insert(&link1).unwrap();
    dao.insert(&link2).unwrap();

    let list_c = Criteria::new();
    match dao.list(&list_c.name(inserted_name)) {
        Ok(v) => assert!(v.len() == 2),
        Err(err) => panic!("LIST FAILED: {}", err),
    }
}
<|endoftext|>


<commit_before>//! On windows we cannot get read (or write) events for sockets in the
//! IOCP model; we only can run asynchronous reads!
//!
//! So we need to use select() to poll for read, and run it in a
//! separate thread.  To cancel that select() we'd need another socket
//! to select() for, and there is no socketpair() - we could use a
//! loopback TCP connection, but a firewall might block it.
//!
//! Instead we use a small (1 second) timeout for the select; it is only
//! used to terminate the thread anyway.
//!
//! This of course wastes one thread per fd we want to watch; a bigger
//! solution would reuse the same backend thread over and over, but then
//! we'd have to try the loopback TCP connection to wake it and fall
//! back to a smaller timeout.

use futures_channel::mpsc as futures_mpsc;
use futures_util::StreamExt;
use log::debug;
use std::{
	io,
	os::raw::c_int,
	sync::{
		mpsc as std_mpsc,
		Mutex,
	},
	task::{
		Context,
		Poll,
	},
	thread,
	time::Duration,
};
use winapi::um::winsock2;

use self::fd_set::FdSet;

pub(crate) struct ReadProcessor {
	fd: libc::c_int,
	poll: PollReadFd,
}

impl ReadProcessor {
	pub(crate) fn new(fd: libc::c_int) -> io::Result<Self> {
		Ok(Self {
			fd,
			poll: PollReadFd::new(fd)?,
		})
	}

	/// call "p" until fd is no longer readable
	pub(crate) fn process<P>(&mut self, cx: &mut Context<'_>, mut p: P) -> io::Result<()>
	where
		P: FnMut() -> io::Result<()>,
	{
		match self.poll.poll_read_ready(cx)? {
			Poll::Ready(()) => {
				while is_readable(self.fd)? {
					p()?;
				}
				self.poll.clear_read_ready(cx)?;
			},
			Poll::Pending => (),
		}
		Ok(())
	}
}

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
enum PollRequest {
	Poll,
	Close,
}

struct SelectFdRead {
	fd: c_int,
	read_fds: FdSet,
}
impl SelectFdRead {
	pub fn new(fd: c_int) -> Self {
		Self {
			fd,
			read_fds: FdSet::new(),
		}
	}

	pub fn select(&mut self, timeout: Option<Duration>) -> bool {
		use std::ptr::null_mut;
		let mut timeout = timeout.map(|timeout| winsock2::timeval {
			tv_sec: timeout.as_secs() as libc::c_long,
			tv_usec: (timeout.subsec_nanos() / 1000) as libc::c_long,
		});
		self.read_fds.set(self.fd);
		unsafe {
			winsock2::select(
				self.fd + 1,
				self.read_fds.inner(),
				null_mut(),
				null_mut(),
				timeout.as_mut().map(|x| x as *mut _).unwrap_or(null_mut()),
			);
		}
		self.read_fds.is_set(self.fd)
	}
}

struct Inner {
	/// file descriptor to watch read events for
	fd: c_int,
	/// background select thread
	_thread: thread::JoinHandle<()>,
	/// either the select thread is running a Poll request or we manually
	/// sent a response through `send_response`
	pending_request: bool,
	/// send poll or close request to select thread
	send_request: std_mpsc::SyncSender<PollRequest>,
	/// when clear_read_ready() is called we use this to trigger a response if
	/// we already know the read event is pending
	send_response: futures_mpsc::Sender<()>,
	/// a response means a read event is pending
	recv_response: futures_mpsc::Receiver<()>,
}

impl Inner {
	fn poll_read_ready(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
		debug!("poll read");
		if !self.pending_request {
			self.pending_request = true;
			let mut read_fds = SelectFdRead::new(self.fd);
			if read_fds.select(Some(Duration::from_millis(0))) {
				debug!("poll read: local ready");
				return Poll::Ready(Ok(()));
			} else {
				debug!("poll read: not ready, start thread");
				self.send_request
					.send(PollRequest::Poll)
					.expect("select thread terminated");
			}
		}

		match self.recv_response.poll_next_unpin(cx) {
			Poll::Ready(None) => unreachable!(), // can't be disconnected
			Poll::Ready(Some(())) => {
				debug!("poll read: thread ready");
				self.pending_request = false;
				Poll::Ready(Ok(()))
			},
			Poll::Pending => {
				debug!("poll read: thread not ready");
				Poll::Pending
			},
		}
	}

	// read() return EAGAIN; re-trigger polling and notify cx
	fn clear_read_ready(&mut self, cx: &mut Context<'_>) -> io::Result<()> {
		// we need to get Poll::Pending from recv_response.poll to make sure `cx` was registered
		match self.recv_response.poll_next_unpin(cx) {
			Poll::Ready(None) => unreachable!(), // can't be disconnected
			Poll::Ready(Some(())) => {
				// was ready. damn...
				assert!(self.pending_request);
				// try again - can't be ready again, but register context
				match self.recv_response.poll_next_unpin(cx) {
					Poll::Ready(None) => unreachable!(),     // can't be disconnected
					Poll::Ready(Some(())) => unreachable!(), // no one could have sent this
					Poll::Pending => (),
				}
				// now send a response - it was ready after all
				self.send_response
					.try_send(())
					.expect("channel can't be full or disconnected");
			},
			Poll::Pending => {
				// yay!
				//
				// already on the way (background thread polling)?
				if !self.pending_request {
					// now we need something to trigger a response
					self.pending_request = true;
					let mut read_fds = SelectFdRead::new(self.fd);
					if read_fds.select(Some(Duration::from_millis(0))) {
						debug!("poll need read: local ready");
						// ready, send a response
						self.send_response
							.try_send(())
							.expect("channel can't be full or disconnected");
					} else {
						debug!("poll need read: not ready, start thread");
						self.send_request
							.send(PollRequest::Poll)
							.expect("select thread terminated");
					}
				}
			},
		}
		Ok(())
	}
}

fn is_readable(fd: c_int) -> io::Result<bool> {
	let mut read_fds = SelectFdRead::new(fd);
	Ok(read_fds.select(Some(Duration::from_millis(0))))
}

struct PollReadFd(Mutex<Inner>);
impl PollReadFd {
	/// does not take overship of fd
	fn new(fd: c_int) -> io::Result<Self> {
		// buffer one request for "Close"
		let (send_request, recv_request) = std_mpsc::sync_channel(1);
		// buffer one notification
		let (mut send_response, recv_response) = futures_mpsc::channel(1);
		let outer_send_response = send_response.clone();

		let thread = thread::spawn(move || {
			let mut read_fds = SelectFdRead::new(fd);
			loop {
				debug!("[select thread] waiting for request");
				match recv_request.recv() {
					Ok(PollRequest::Poll) => (),
					Ok(PollRequest::Close) => return,
					Err(_) => return,
				}
				debug!("[select thread] start polling");

				while !read_fds.select(Some(Duration::from_millis(1000))) {
					match recv_request.try_recv() {
						Ok(PollRequest::Poll) => unreachable!(),
						Ok(PollRequest::Close) => return,
						Err(_) => (), // back to select()
					}
				}

				debug!("[select thread] read event");

				// Can only fail if the other end is dropped
				let _ = futures_executor::block_on(send_response.send(()));
			}
		});

		Ok(Self(Mutex::new(Inner {
			fd,
			_thread: thread,
			pending_request: false,
			send_request,
			send_response: outer_send_response,
			recv_response,
		})))
	}

	fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
		self.0.lock().expect("mutex poisoned").poll_read_ready(cx)
	}

	fn clear_read_ready(&self, cx: &mut Context<'_>) -> io::Result<()> {
		self.0.lock().expect("mutex poisoned").clear_read_ready(cx)
	}
}

impl Drop for PollReadFd {
	fn drop(&mut self) {
		let _ = self
			.0
			.get_mut()
			.expect("mutex poisoned")
			.send_request
			.send(PollRequest::Close);
	}
}

mod fd_set {
	use libc::{
		c_int,
		c_uint,
	};
	use std::{
		mem::MaybeUninit,
		ptr,
	};
	use winapi::um::winsock2::{
		fd_set,
		u_int,
		FD_SETSIZE,
		SOCKET,
	};

	/// Layout compatible struct of `fd_set`, but it holds maybe uninitialized `fd_array`.
	///
	/// # Safety
	/// The first `fd_count` slots of `fd_array` must be initialized,
	/// and the rest may be uninitialized.
	#[repr(C)]
	pub(super) struct FdSet {
		fd_count: u_int,
		fd_array: [MaybeUninit<SOCKET>; FD_SETSIZE],
	}

	impl FdSet {
		pub fn new() -> Self {
			Self {
				fd_count: 0,
				// safe according to:
				// https://doc.rust-lang.org/std/mem/union.MaybeUninit.html#initializing-an-array-element-by-element
				fd_array: unsafe { MaybeUninit::uninit().assume_init() },
			}
		}

		pub fn inner(&mut self) -> *mut fd_set {
			self as *mut Self as *mut _
		}

		pub fn set(&mut self, fd: c_int) {
			if self.is_set(fd) {
				return;
			}
			let count = self.fd_count as usize;
			if count < FD_SETSIZE {
				let fd = fd as c_uint as SOCKET;
				// This is safe because this slot is uninitialized.
				unsafe { ptr::write(self.fd_array[count].as_mut_ptr(), fd) };
				self.fd_count += 1;
			}
		}

		pub fn is_set(&self, fd: c_int) -> bool {
			let fd = fd as c_uint as SOCKET;
			let count = self.fd_count as usize;
			self.fd_array[..count].iter().any(|slot| {
				// This is safe because it's reading from first `fd_count` slots.
				fd == unsafe { ptr::read(slot.as_ptr()) }
			})
		}
	}

	#[cfg(test)]
	mod tests {
		use super::*;
		use std::mem::{
			needs_drop,
			transmute,
		};

		// Check that `FdSet` is layout compatible with `fd_set`.
		#[test]
		fn fd_set_layout_compatible() {
			let mut fd_set = FdSet::new();
			(0..FD_SETSIZE as c_int).for_each(|i| fd_set.set(i));
			let fd_set: fd_set = unsafe { transmute(fd_set) };
			assert_eq!(fd_set.fd_count, FD_SETSIZE as u32);
			for i in 0..FD_SETSIZE as usize {
				assert_eq!(fd_set.fd_array[i], i as SOCKET);
			}
		}

		// Check that `SOCKET` doesn't need to be dropped, so that we don't need to
		// implement `Drop` for `FdSet`.
		#[test]
		fn socket_not_drop() {
			assert!(!needs_drop::<SOCKET>());
		}
	}
}
<commit_msg>Added futures_util::SinkExt<commit_after>//! On windows we cannot get read (or write) events for sockets in the
//! IOCP model; we only can run asynchronous reads!
//!
//! So we need to use select() to poll for read, and run it in a
//! separate thread.  To cancel that select() we'd need another socket
//! to select() for, and there is no socketpair() - we could use a
//! loopback TCP connection, but a firewall might block it.
//!
//! Instead we use a small (1 second) timeout for the select; it is only
//! used to terminate the thread anyway.
//!
//! This of course wastes one thread per fd we want to watch; a bigger
//! solution would reuse the same backend thread over and over, but then
//! we'd have to try the loopback TCP connection to wake it and fall
//! back to a smaller timeout.

use futures_channel::mpsc as futures_mpsc;
use futures_util::{
	SinkExt,
	StreamExt,
};
use log::debug;
use std::{
	io,
	os::raw::c_int,
	sync::{
		mpsc as std_mpsc,
		Mutex,
	},
	task::{
		Context,
		Poll,
	},
	thread,
	time::Duration,
};
use winapi::um::winsock2;

use self::fd_set::FdSet;

pub(crate) struct ReadProcessor {
	fd: libc::c_int,
	poll: PollReadFd,
}

impl ReadProcessor {
	pub(crate) fn new(fd: libc::c_int) -> io::Result<Self> {
		Ok(Self {
			fd,
			poll: PollReadFd::new(fd)?,
		})
	}

	/// call "p" until fd is no longer readable
	pub(crate) fn process<P>(&mut self, cx: &mut Context<'_>, mut p: P) -> io::Result<()>
	where
		P: FnMut() -> io::Result<()>,
	{
		match self.poll.poll_read_ready(cx)? {
			Poll::Ready(()) => {
				while is_readable(self.fd)? {
					p()?;
				}
				self.poll.clear_read_ready(cx)?;
			},
			Poll::Pending => (),
		}
		Ok(())
	}
}

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
enum PollRequest {
	Poll,
	Close,
}

struct SelectFdRead {
	fd: c_int,
	read_fds: FdSet,
}
impl SelectFdRead {
	pub fn new(fd: c_int) -> Self {
		Self {
			fd,
			read_fds: FdSet::new(),
		}
	}

	pub fn select(&mut self, timeout: Option<Duration>) -> bool {
		use std::ptr::null_mut;
		let mut timeout = timeout.map(|timeout| winsock2::timeval {
			tv_sec: timeout.as_secs() as libc::c_long,
			tv_usec: (timeout.subsec_nanos() / 1000) as libc::c_long,
		});
		self.read_fds.set(self.fd);
		unsafe {
			winsock2::select(
				self.fd + 1,
				self.read_fds.inner(),
				null_mut(),
				null_mut(),
				timeout.as_mut().map(|x| x as *mut _).unwrap_or(null_mut()),
			);
		}
		self.read_fds.is_set(self.fd)
	}
}

struct Inner {
	/// file descriptor to watch read events for
	fd: c_int,
	/// background select thread
	_thread: thread::JoinHandle<()>,
	/// either the select thread is running a Poll request or we manually
	/// sent a response through `send_response`
	pending_request: bool,
	/// send poll or close request to select thread
	send_request: std_mpsc::SyncSender<PollRequest>,
	/// when clear_read_ready() is called we use this to trigger a response if
	/// we already know the read event is pending
	send_response: futures_mpsc::Sender<()>,
	/// a response means a read event is pending
	recv_response: futures_mpsc::Receiver<()>,
}

impl Inner {
	fn poll_read_ready(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
		debug!("poll read");
		if !self.pending_request {
			self.pending_request = true;
			let mut read_fds = SelectFdRead::new(self.fd);
			if read_fds.select(Some(Duration::from_millis(0))) {
				debug!("poll read: local ready");
				return Poll::Ready(Ok(()));
			} else {
				debug!("poll read: not ready, start thread");
				self.send_request
					.send(PollRequest::Poll)
					.expect("select thread terminated");
			}
		}

		match self.recv_response.poll_next_unpin(cx) {
			Poll::Ready(None) => unreachable!(), // can't be disconnected
			Poll::Ready(Some(())) => {
				debug!("poll read: thread ready");
				self.pending_request = false;
				Poll::Ready(Ok(()))
			},
			Poll::Pending => {
				debug!("poll read: thread not ready");
				Poll::Pending
			},
		}
	}

	// read() return EAGAIN; re-trigger polling and notify cx
	fn clear_read_ready(&mut self, cx: &mut Context<'_>) -> io::Result<()> {
		// we need to get Poll::Pending from recv_response.poll to make sure `cx` was registered
		match self.recv_response.poll_next_unpin(cx) {
			Poll::Ready(None) => unreachable!(), // can't be disconnected
			Poll::Ready(Some(())) => {
				// was ready. damn...
				assert!(self.pending_request);
				// try again - can't be ready again, but register context
				match self.recv_response.poll_next_unpin(cx) {
					Poll::Ready(None) => unreachable!(),     // can't be disconnected
					Poll::Ready(Some(())) => unreachable!(), // no one could have sent this
					Poll::Pending => (),
				}
				// now send a response - it was ready after all
				self.send_response
					.try_send(())
					.expect("channel can't be full or disconnected");
			},
			Poll::Pending => {
				// yay!
				//
				// already on the way (background thread polling)?
				if !self.pending_request {
					// now we need something to trigger a response
					self.pending_request = true;
					let mut read_fds = SelectFdRead::new(self.fd);
					if read_fds.select(Some(Duration::from_millis(0))) {
						debug!("poll need read: local ready");
						// ready, send a response
						self.send_response
							.try_send(())
							.expect("channel can't be full or disconnected");
					} else {
						debug!("poll need read: not ready, start thread");
						self.send_request
							.send(PollRequest::Poll)
							.expect("select thread terminated");
					}
				}
			},
		}
		Ok(())
	}
}

fn is_readable(fd: c_int) -> io::Result<bool> {
	let mut read_fds = SelectFdRead::new(fd);
	Ok(read_fds.select(Some(Duration::from_millis(0))))
}

struct PollReadFd(Mutex<Inner>);
impl PollReadFd {
	/// does not take overship of fd
	fn new(fd: c_int) -> io::Result<Self> {
		// buffer one request for "Close"
		let (send_request, recv_request) = std_mpsc::sync_channel(1);
		// buffer one notification
		let (mut send_response, recv_response) = futures_mpsc::channel(1);
		let outer_send_response = send_response.clone();

		let thread = thread::spawn(move || {
			let mut read_fds = SelectFdRead::new(fd);
			loop {
				debug!("[select thread] waiting for request");
				match recv_request.recv() {
					Ok(PollRequest::Poll) => (),
					Ok(PollRequest::Close) => return,
					Err(_) => return,
				}
				debug!("[select thread] start polling");

				while !read_fds.select(Some(Duration::from_millis(1000))) {
					match recv_request.try_recv() {
						Ok(PollRequest::Poll) => unreachable!(),
						Ok(PollRequest::Close) => return,
						Err(_) => (), // back to select()
					}
				}

				debug!("[select thread] read event");

				// Can only fail if the other end is dropped
				let _ = futures_executor::block_on(send_response.send(()));
			}
		});

		Ok(Self(Mutex::new(Inner {
			fd,
			_thread: thread,
			pending_request: false,
			send_request,
			send_response: outer_send_response,
			recv_response,
		})))
	}

	fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
		self.0.lock().expect("mutex poisoned").poll_read_ready(cx)
	}

	fn clear_read_ready(&self, cx: &mut Context<'_>) -> io::Result<()> {
		self.0.lock().expect("mutex poisoned").clear_read_ready(cx)
	}
}

impl Drop for PollReadFd {
	fn drop(&mut self) {
		let _ = self
			.0
			.get_mut()
			.expect("mutex poisoned")
			.send_request
			.send(PollRequest::Close);
	}
}

mod fd_set {
	use libc::{
		c_int,
		c_uint,
	};
	use std::{
		mem::MaybeUninit,
		ptr,
	};
	use winapi::um::winsock2::{
		fd_set,
		u_int,
		FD_SETSIZE,
		SOCKET,
	};

	/// Layout compatible struct of `fd_set`, but it holds maybe uninitialized `fd_array`.
	///
	/// # Safety
	/// The first `fd_count` slots of `fd_array` must be initialized,
	/// and the rest may be uninitialized.
	#[repr(C)]
	pub(super) struct FdSet {
		fd_count: u_int,
		fd_array: [MaybeUninit<SOCKET>; FD_SETSIZE],
	}

	impl FdSet {
		pub fn new() -> Self {
			Self {
				fd_count: 0,
				// safe according to:
				// https://doc.rust-lang.org/std/mem/union.MaybeUninit.html#initializing-an-array-element-by-element
				fd_array: unsafe { MaybeUninit::uninit().assume_init() },
			}
		}

		pub fn inner(&mut self) -> *mut fd_set {
			self as *mut Self as *mut _
		}

		pub fn set(&mut self, fd: c_int) {
			if self.is_set(fd) {
				return;
			}
			let count = self.fd_count as usize;
			if count < FD_SETSIZE {
				let fd = fd as c_uint as SOCKET;
				// This is safe because this slot is uninitialized.
				unsafe { ptr::write(self.fd_array[count].as_mut_ptr(), fd) };
				self.fd_count += 1;
			}
		}

		pub fn is_set(&self, fd: c_int) -> bool {
			let fd = fd as c_uint as SOCKET;
			let count = self.fd_count as usize;
			self.fd_array[..count].iter().any(|slot| {
				// This is safe because it's reading from first `fd_count` slots.
				fd == unsafe { ptr::read(slot.as_ptr()) }
			})
		}
	}

	#[cfg(test)]
	mod tests {
		use super::*;
		use std::mem::{
			needs_drop,
			transmute,
		};

		// Check that `FdSet` is layout compatible with `fd_set`.
		#[test]
		fn fd_set_layout_compatible() {
			let mut fd_set = FdSet::new();
			(0..FD_SETSIZE as c_int).for_each(|i| fd_set.set(i));
			let fd_set: fd_set = unsafe { transmute(fd_set) };
			assert_eq!(fd_set.fd_count, FD_SETSIZE as u32);
			for i in 0..FD_SETSIZE as usize {
				assert_eq!(fd_set.fd_array[i], i as SOCKET);
			}
		}

		// Check that `SOCKET` doesn't need to be dropped, so that we don't need to
		// implement `Drop` for `FdSet`.
		#[test]
		fn socket_not_drop() {
			assert!(!needs_drop::<SOCKET>());
		}
	}
}
<|endoftext|>