<commit_before>/*
 * Copyright (c) 2013, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

use std;
use common::*;
use arena::*;
use layout;

pub struct ReaderOptions {
    traversalLimitInWords : u64,
    nestingLimit : uint
}

pub static DEFAULT_READER_OPTIONS : ReaderOptions =
    ReaderOptions { traversalLimitInWords : 8 * 1024 * 1024, nestingLimit : 64 };

pub struct MessageReader<'self> {
    segments : &'self [ &'self [u8]],
    options : ReaderOptions,
//    arena : ReaderArena<'self>
}

type SegmentId = u32;

impl <'self> MessageReader<'self> {

    #[inline]
    pub fn getSegment(&self, id : uint) -> &'self [u8] {
        self.segments[id]
    }

    #[inline]
    pub fn getSegmentReader<'a>(&'a self, id : SegmentId) -> SegmentReader<'a> {
        SegmentReader { messageReader : self, segment : self.getSegment(id as uint) }
    }

    #[inline]
    pub fn getOptions<'a>(&'a self) -> &'a ReaderOptions {
        return &self.options;
    }

    pub fn getRoot<'a>(&'a self) -> layout::StructReader<'a> {
        let segment = self.getSegmentReader(0);

        return layout::StructReader::readRoot(0, segment,
                                              self.options.nestingLimit as int);
    }

}

pub enum AllocationStrategy {
    FIXED_SIZE,
    GROW_HEURISTICALLY
}

pub static SUGGESTED_FIRST_SEGMENT_WORDS : uint = 1024;
pub static SUGGESTED_ALLOCATION_STRATEGY : AllocationStrategy = GROW_HEURISTICALLY;

pub struct MessageBuilder {
    nextSize : uint,
    allocationStrategy : AllocationStrategy,
    segmentBuilders : ~[~SegmentBuilder],

    segments : ~[~[u8]]
    // It would probably be nicer if this were a vector of @mut[u8]s
    // and the SegmentBuilders also directly held their respective
    // @mut[u8]s. Only problem is, I don't know how to implement
    // `asReader` in that case.
}

impl MessageBuilder {

    pub fn new(firstSegmentWords : uint, allocationStrategy : AllocationStrategy)
        -> ~MessageBuilder {
        let mut result = ~MessageBuilder {
            nextSize : firstSegmentWords,
            allocationStrategy : allocationStrategy,
            segmentBuilders : ~[],
            segments : ~[]
        };

        let builder =
            ~SegmentBuilder::new(std::ptr::to_mut_unsafe_ptr(result), firstSegmentWords);

        result.segments.push(allocate_zeroed_bytes(firstSegmentWords * BYTES_PER_WORD));
        result.segmentBuilders.push(builder);

        result
    }

    pub fn new_default() -> ~MessageBuilder {
        MessageBuilder::new(SUGGESTED_FIRST_SEGMENT_WORDS, SUGGESTED_ALLOCATION_STRATEGY)
    }

    pub fn allocateSegment(&mut self, minimumSize : WordCount) -> *mut SegmentBuilder {
        let size = std::cmp::max(minimumSize, self.nextSize);
        let segment = allocate_zeroed_bytes(size * BYTES_PER_WORD);
        let mut result = ~SegmentBuilder::new(self, size);
        let result_ptr = std::ptr::to_mut_unsafe_ptr(result);
        self.segments.push(segment);
        self.segmentBuilders.push(result);

        match self.allocationStrategy {
            GROW_HEURISTICALLY => { self.nextSize += size; }
            _ => { }
        }

        result_ptr
    }

    pub fn getSegmentWithAvailable(&mut self, minimumAvailable : WordCount)
        -> *mut SegmentBuilder {
        if (self.segmentBuilders.last().available() >= minimumAvailable) {
            return std::ptr::to_mut_unsafe_ptr(self.segmentBuilders[self.segments.len() - 1]);
        } else {
            return self.allocateSegment(minimumAvailable);
        }
    }

    pub fn initRoot<T : layout::HasStructSize + layout::FromStructBuilder>(&mut self) -> T {
        // Rolled in this stuff form getRootSegment.
        let rootSegment = std::ptr::to_mut_unsafe_ptr(self.segmentBuilders[0]);

        let unused_self : Option<T> = None;

        match self.segmentBuilders[0].allocate(WORDS_PER_POINTER) {
            None => {fail!("could not allocate root pointer") }
            Some(location) => {
                //assert!(location == 0,
                //        "First allocated word of new segment was not at offset 0");

                let sb = layout::StructBuilder::initRoot(
                    rootSegment,
                    unsafe {std::cast::transmute(location)},
                    layout::HasStructSize::structSize(unused_self));

                return layout::FromStructBuilder::fromStructBuilder(sb);
            }
        }
    }

    pub fn asReader<T>(& self, f : &fn(r : MessageReader) -> T) -> T {
        let mut segments : ~[&[u8]] = ~[];

        for ii in range(0, self.segments.len()) {
            segments.push(self.segments[ii].as_slice());
        }

        f(MessageReader {segments : segments, options : DEFAULT_READER_OPTIONS})
    }

}
<commit_msg>remove obsolete comment<commit_after>/*
 * Copyright (c) 2013, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

use std;
use common::*;
use arena::*;
use layout;

pub struct ReaderOptions {
    traversalLimitInWords : u64,
    nestingLimit : uint
}

pub static DEFAULT_READER_OPTIONS : ReaderOptions =
    ReaderOptions { traversalLimitInWords : 8 * 1024 * 1024, nestingLimit : 64 };

pub struct MessageReader<'self> {
    segments : &'self [ &'self [u8]],
    options : ReaderOptions,
//    arena : ReaderArena<'self>
}

type SegmentId = u32;

impl <'self> MessageReader<'self> {

    #[inline]
    pub fn getSegment(&self, id : uint) -> &'self [u8] {
        self.segments[id]
    }

    #[inline]
    pub fn getSegmentReader<'a>(&'a self, id : SegmentId) -> SegmentReader<'a> {
        SegmentReader { messageReader : self, segment : self.getSegment(id as uint) }
    }

    #[inline]
    pub fn getOptions<'a>(&'a self) -> &'a ReaderOptions {
        return &self.options;
    }

    pub fn getRoot<'a>(&'a self) -> layout::StructReader<'a> {
        let segment = self.getSegmentReader(0);

        return layout::StructReader::readRoot(0, segment,
                                              self.options.nestingLimit as int);
    }

}

pub enum AllocationStrategy {
    FIXED_SIZE,
    GROW_HEURISTICALLY
}

pub static SUGGESTED_FIRST_SEGMENT_WORDS : uint = 1024;
pub static SUGGESTED_ALLOCATION_STRATEGY : AllocationStrategy = GROW_HEURISTICALLY;

pub struct MessageBuilder {
    nextSize : uint,
    allocationStrategy : AllocationStrategy,
    segmentBuilders : ~[~SegmentBuilder],
    segments : ~[~[u8]]
}

impl MessageBuilder {

    pub fn new(firstSegmentWords : uint, allocationStrategy : AllocationStrategy)
        -> ~MessageBuilder {
        let mut result = ~MessageBuilder {
            nextSize : firstSegmentWords,
            allocationStrategy : allocationStrategy,
            segmentBuilders : ~[],
            segments : ~[]
        };

        let builder =
            ~SegmentBuilder::new(std::ptr::to_mut_unsafe_ptr(result), firstSegmentWords);

        result.segments.push(allocate_zeroed_bytes(firstSegmentWords * BYTES_PER_WORD));
        result.segmentBuilders.push(builder);

        result
    }

    pub fn new_default() -> ~MessageBuilder {
        MessageBuilder::new(SUGGESTED_FIRST_SEGMENT_WORDS, SUGGESTED_ALLOCATION_STRATEGY)
    }

    pub fn allocateSegment(&mut self, minimumSize : WordCount) -> *mut SegmentBuilder {
        let size = std::cmp::max(minimumSize, self.nextSize);
        let segment = allocate_zeroed_bytes(size * BYTES_PER_WORD);
        let mut result = ~SegmentBuilder::new(self, size);
        let result_ptr = std::ptr::to_mut_unsafe_ptr(result);
        self.segments.push(segment);
        self.segmentBuilders.push(result);

        match self.allocationStrategy {
            GROW_HEURISTICALLY => { self.nextSize += size; }
            _ => { }
        }

        result_ptr
    }

    pub fn getSegmentWithAvailable(&mut self, minimumAvailable : WordCount)
        -> *mut SegmentBuilder {
        if (self.segmentBuilders.last().available() >= minimumAvailable) {
            return std::ptr::to_mut_unsafe_ptr(self.segmentBuilders[self.segments.len() - 1]);
        } else {
            return self.allocateSegment(minimumAvailable);
        }
    }

    pub fn initRoot<T : layout::HasStructSize + layout::FromStructBuilder>(&mut self) -> T {
        // Rolled in this stuff form getRootSegment.
        let rootSegment = std::ptr::to_mut_unsafe_ptr(self.segmentBuilders[0]);

        let unused_self : Option<T> = None;

        match self.segmentBuilders[0].allocate(WORDS_PER_POINTER) {
            None => {fail!("could not allocate root pointer") }
            Some(location) => {
                //assert!(location == 0,
                //        "First allocated word of new segment was not at offset 0");

                let sb = layout::StructBuilder::initRoot(
                    rootSegment,
                    unsafe {std::cast::transmute(location)},
                    layout::HasStructSize::structSize(unused_self));

                return layout::FromStructBuilder::fromStructBuilder(sb);
            }
        }
    }

    pub fn asReader<T>(& self, f : &fn(r : MessageReader) -> T) -> T {
        let mut segments : ~[&[u8]] = ~[];

        for ii in range(0, self.segments.len()) {
            segments.push(self.segments[ii].as_slice());
        }

        f(MessageReader {segments : segments, options : DEFAULT_READER_OPTIONS})
    }

}
<|endoftext|>


<commit_before>use core::ptr::{Unique};

use arch::{kernel_start_paddr, kernel_start_vaddr,
           kernel_end_paddr, kernel_end_vaddr};
use arch::paging::{PTEntry, PML4, PDPT, PD, PT,
                   pml4_index, pdpt_index, pd_index, pt_index,
                   BASE_PAGE_LENGTH, LARGE_PAGE_LENGTH};
use arch::interrupt::{APICPage};
use arch::{KERNEL_BASE};
use common::{PAddr, VAddr, MemoryRegion};
use util::{block_count, align_up, ExternReadonlyObject, ExternMutex};
use core::ops::{Deref};

use arch::addr;

extern {
    static mut init_pd: PD;
    static kernel_stack_guard_page: u64;
}

// Below should be used BEFORE switching to new page table structure.
const INITIAL_ALLOC_START_VADDR: VAddr = VAddr::new(KERNEL_BASE + 0xc00000);
const INITIAL_ALLOC_PML4_OFFSET: usize = 0x0000;
const INITIAL_ALLOC_PDPT_OFFSET: usize = 0x1000;
const INITIAL_ALLOC_PD_OFFSET: usize = 0x2000;
const INITIAL_ALLOC_OBJECT_POOL_PT_OFFSET: usize = 0x3000;
const INITIAL_ALLOC_KERNEL_PT_START_OFFSET: usize = 0x4000;

// Below should be used AFTER switching to new page table structure.
pub const OBJECT_POOL_START_VADDR: VAddr = VAddr::new(KERNEL_BASE + 0xe00000);
pub const OBJECT_POOL_SIZE: usize = 511;
pub const OBJECT_POOL_PT_VADDR: VAddr = VAddr::new(KERNEL_BASE + 0xfff000);

pub static APIC_PAGE: ExternMutex<APICPage> = unsafe { ExternMutex::new(None) };

// Variables
static INITIAL_PD: ExternMutex<PD> =
    unsafe { ExternMutex::new(Some(&init_pd as *const _)) };

pub static OBJECT_POOL_PT: ExternMutex<[PTEntry; OBJECT_POOL_SIZE]> =
    unsafe { ExternMutex::new(None) };
pub static KERNEL_PML4: ExternReadonlyObject<PML4> =
    unsafe { ExternReadonlyObject::new() };
pub static KERNEL_PDPT: ExternReadonlyObject<PDPT> =
    unsafe { ExternReadonlyObject::new() };
pub static KERNEL_PD: ExternReadonlyObject<PD> =
    unsafe { ExternReadonlyObject::new() };

fn kernel_stack_guard_page_vaddr() -> VAddr {
    VAddr::from((&kernel_stack_guard_page as *const _) as u64)
}

fn alloc_kernel_pml4(region: &mut MemoryRegion, alloc_base: PAddr) -> Unique<PML4> {
    use arch::paging::PML4Entry;

    let paddr = alloc_base + INITIAL_ALLOC_PML4_OFFSET;
    let vaddr = INITIAL_ALLOC_START_VADDR + INITIAL_ALLOC_PML4_OFFSET;

    log!("pml4, paddr: 0x{:x}, vaddr: 0x{:x}", paddr, vaddr);

    let mut pml4_unique = unsafe { Unique::new(vaddr.into(): usize as *mut PML4) };

    {
        let mut pml4 = unsafe { pml4_unique.get_mut() };
        *pml4 = [PML4Entry::empty(); 512];
    }

    region.move_up(paddr + BASE_PAGE_LENGTH);

    unsafe { KERNEL_PML4.bootstrap(*pml4_unique.deref(), paddr); }

    pml4_unique
}

fn alloc_kernel_pdpt(region: &mut MemoryRegion, pml4: &mut PML4, alloc_base: PAddr) -> Unique<PDPT> {
    use arch::paging::{PDPTEntry, PML4Entry, PML4_P, PML4_RW};

    let paddr = alloc_base + INITIAL_ALLOC_PDPT_OFFSET;
    let vaddr = INITIAL_ALLOC_START_VADDR + INITIAL_ALLOC_PDPT_OFFSET;

    log!("pdpt, paddr: 0x{:x}, vaddr: 0x{:x}", paddr, vaddr);

    let mut pdpt_unique = unsafe { Unique::new(vaddr.into(): usize as *mut PDPT) };

    {
        let mut pdpt = unsafe { pdpt_unique.get_mut() };
        *pdpt = [PDPTEntry::empty(); 512];
    }

    region.move_up(paddr + BASE_PAGE_LENGTH);

    pml4[pml4_index(VAddr::from(KERNEL_BASE))] = PML4Entry::new(paddr, PML4_P | PML4_RW);

    unsafe { KERNEL_PDPT.bootstrap(*pdpt_unique.deref(), paddr) }

    pdpt_unique
}

fn alloc_kernel_pd(region: &mut MemoryRegion, pdpt: &mut PDPT, alloc_base: PAddr) -> Unique<PD> {
    use arch::paging::{PDEntry, PDPTEntry, PDPT_P, PDPT_RW};

    let paddr = alloc_base + INITIAL_ALLOC_PD_OFFSET;
    let vaddr = INITIAL_ALLOC_START_VADDR + INITIAL_ALLOC_PD_OFFSET;

    log!("pd, paddr: 0x{:x}, vaddr: 0x{:x}", paddr, vaddr);

    let mut pd_unique = unsafe { Unique::new(vaddr.into(): usize as *mut PD) };

    {
        let mut pd = unsafe { pd_unique.get_mut() };
        *pd = [PDEntry::empty(); 512];
    }

    region.move_up(paddr + BASE_PAGE_LENGTH);

    pdpt[pdpt_index(VAddr::from(KERNEL_BASE))] = PDPTEntry::new(paddr, PDPT_P | PDPT_RW);

    unsafe { KERNEL_PD.bootstrap(*pd_unique.deref(), paddr); }

    pd_unique
}

fn alloc_object_pool_pt(region: &mut MemoryRegion, pd: &mut PD, alloc_base: PAddr) -> Unique<PT> {
    use arch::paging::{PTEntry, PDEntry, PD_P, PD_RW, PT_P, PT_RW};

    let paddr = alloc_base + INITIAL_ALLOC_OBJECT_POOL_PT_OFFSET;
    let vaddr = INITIAL_ALLOC_START_VADDR + INITIAL_ALLOC_OBJECT_POOL_PT_OFFSET;

    log!("object_pool_pt, paddr: 0x{:x}, vaddr: 0x{:x}", paddr, vaddr);

    let mut pt_unique = unsafe { Unique::new(vaddr.into(): usize as *mut PT) };

    {
        let mut pt = unsafe { pt_unique.get_mut() };
        *pt = [PTEntry::empty(); 512];

        let reverse_pt_index = pt_index(OBJECT_POOL_PT_VADDR);
        log!("reverse_pt_index: {}", reverse_pt_index);

        pt[reverse_pt_index] = PTEntry::new(paddr, PT_P | PT_RW);
    }

    region.move_up(paddr + BASE_PAGE_LENGTH);

    pd[pd_index(OBJECT_POOL_START_VADDR)] = PDEntry::new(paddr, PD_P | PD_RW);

    pt_unique
}

fn alloc_apic(pt: &mut PT, offset_size: usize, alloc_base: PAddr) {
    use arch::paging::{PT_P, PT_RW, PT_PWT, PT_PCD};
    use x86::shared::msr;
    use core::mem;

    let apic_msr = unsafe { msr::rdmsr(0x1B) };
    let apic_base = (apic_msr >> 12) << 3;
    assert!(apic_msr & (0x1 << 11) == (0x1 << 11));
    log!("apic base address: 0x{:x}", apic_base);

    assert!(mem::size_of::<APICPage>() == 1024);

    let paddr = PAddr::from(apic_base);
    let vaddr = kernel_start_vaddr() + (offset_size * BASE_PAGE_LENGTH);

    log!("apic allocated at 0x{:x}", vaddr);

    pt[pt_index(vaddr)] = PTEntry::new(paddr, PT_P | PT_RW | PT_PWT | PT_PCD);

    unsafe { APIC_PAGE.bootstrap(vaddr.into(): usize as *const APICPage) };
}

fn alloc_kernel_page(pt: &mut PT, offset_size: usize, alloc_base: PAddr) {
    use arch::paging::{PT_P, PT_RW};

    let paddr = kernel_start_paddr() + (offset_size * BASE_PAGE_LENGTH);
    let vaddr = kernel_start_vaddr() + (offset_size * BASE_PAGE_LENGTH);

    log!("kernel page allocated at 0x{:x}", vaddr);

    pt[pt_index(vaddr)] = PTEntry::new(paddr, PT_P | PT_RW);
}

fn alloc_kernel_guard_page(pt: &mut PT, offset_size: usize, alloc_base: PAddr) {
    use arch::paging::{PT_P, PT_RW};

    let paddr = kernel_start_paddr() + (offset_size * BASE_PAGE_LENGTH);
    let vaddr = kernel_start_vaddr() + (offset_size * BASE_PAGE_LENGTH);

    log!("guard page allocated at 0x{:x}", vaddr);

    pt[pt_index(vaddr)] = PTEntry::empty();
}

fn alloc_kernel_pts(region: &mut MemoryRegion, pd: &mut PD, alloc_base: PAddr) {
    use arch::paging::{PDEntry, PD_P, PD_RW};

    let kernel_page_size = block_count(kernel_end_paddr().into(): usize -
                                       kernel_start_paddr().into(): usize, BASE_PAGE_LENGTH);
    let guard_page_index = (kernel_stack_guard_page_vaddr().into(): usize -
                            kernel_start_vaddr().into(): usize) / BASE_PAGE_LENGTH;

    log!("guard_page_index: {}", guard_page_index);

    for i in 0..(kernel_page_size + 1) {
        if i % 512 == 0 {
            pd[pd_index(kernel_start_vaddr() + i * BASE_PAGE_LENGTH)] = PDEntry::new(region.start_paddr(), PD_P | PD_RW);
            let npaddr = region.start_paddr() + BASE_PAGE_LENGTH;
            region.move_up(npaddr);
        }

        let pt_entry = pd[pd_index(kernel_start_vaddr() + i * BASE_PAGE_LENGTH)];

        let offset = pt_entry.get_address().into(): usize -
            alloc_base.into(): usize;
        let mut pt_unique = unsafe {
            Unique::new((INITIAL_ALLOC_START_VADDR + offset).into(): usize as *mut PT) };

        if i == guard_page_index {
            alloc_kernel_guard_page(unsafe { pt_unique.get_mut() }, i % 512, alloc_base);
        } else if i == kernel_page_size {
            alloc_apic(unsafe { pt_unique.get_mut() }, i % 512, alloc_base);
        } else {
            alloc_kernel_page(unsafe { pt_unique.get_mut() }, i % 512, alloc_base);
        }
    }
}

// This maps 2MB for allocation region
fn map_alloc_region(alloc_region: &mut MemoryRegion) -> PAddr {
    use arch::paging::{PD_P, PD_RW, PD_PS, PDEntry, flush_all};

    let map_alloc_start_vaddr = INITIAL_ALLOC_START_VADDR;
    let map_alloc_pd_index = pd_index(map_alloc_start_vaddr);
    let map_alloc_start_paddr = align_up(alloc_region.start_paddr(), LARGE_PAGE_LENGTH);

    let mut initial_pd = INITIAL_PD.lock();
    initial_pd[map_alloc_pd_index] = PDEntry::new(map_alloc_start_paddr, PD_P | PD_RW | PD_PS);

    // unsafe { super::paging::flush(map_alloc_start_vaddr); }
    unsafe { flush_all(); }

    alloc_region.move_up(map_alloc_start_paddr);

    map_alloc_start_paddr
}

pub fn init(mut alloc_region: &mut MemoryRegion) {
    use arch::paging::{switch_to};

    let kernel_page_size = block_count(kernel_end_paddr().into(): usize -
                                       kernel_start_paddr().into(): usize, BASE_PAGE_LENGTH);
    let alloc_size = 3 /* PML4, PDPT, and PD, and object pool PT */ +
        block_count(kernel_page_size, 512) /* Kernel page mapping PT */ ;

    assert!(alloc_size <= 512);
    assert!(alloc_size * BASE_PAGE_LENGTH < alloc_region.length());

    // Before allocation, we need to make sure PAddr + alloc_size is mapped.
    // This is done in the init_pd.

    let alloc_base_paddr = map_alloc_region(&mut alloc_region);

    log!("alloc_base_paddr: 0x{:x}", alloc_base_paddr);

    let mut pml4_unique = alloc_kernel_pml4(&mut alloc_region,
                                            alloc_base_paddr);
    let mut pdpt_unique = alloc_kernel_pdpt(&mut alloc_region,
                                            unsafe { pml4_unique.get_mut() },
                                            alloc_base_paddr);
    let mut pd_unique = alloc_kernel_pd(&mut alloc_region,
                                        unsafe { pdpt_unique.get_mut() },
                                        alloc_base_paddr);
    let mut object_pool_pt_unique = alloc_object_pool_pt(&mut alloc_region,
                                                         unsafe { pd_unique.get_mut() },
                                                         alloc_base_paddr);

    alloc_kernel_pts(&mut alloc_region, unsafe { pd_unique.get_mut() }, alloc_base_paddr);

    unsafe {
        INITIAL_PD.unbootstrap();
    }
    unsafe { switch_to(KERNEL_PML4.paddr()); }
    unsafe {
        OBJECT_POOL_PT.bootstrap(OBJECT_POOL_PT_VADDR.into(): usize as *mut _);
    }
}
<commit_msg>Fix APIC base calculation<commit_after>use core::ptr::{Unique};

use arch::{kernel_start_paddr, kernel_start_vaddr,
           kernel_end_paddr, kernel_end_vaddr};
use arch::paging::{PTEntry, PML4, PDPT, PD, PT,
                   pml4_index, pdpt_index, pd_index, pt_index,
                   BASE_PAGE_LENGTH, LARGE_PAGE_LENGTH};
use arch::interrupt::{APICPage};
use arch::{KERNEL_BASE};
use common::{PAddr, VAddr, MemoryRegion};
use util::{block_count, align_up, ExternReadonlyObject, ExternMutex};
use core::ops::{Deref};

use arch::addr;

extern {
    static mut init_pd: PD;
    static kernel_stack_guard_page: u64;
}

// Below should be used BEFORE switching to new page table structure.
const INITIAL_ALLOC_START_VADDR: VAddr = VAddr::new(KERNEL_BASE + 0xc00000);
const INITIAL_ALLOC_PML4_OFFSET: usize = 0x0000;
const INITIAL_ALLOC_PDPT_OFFSET: usize = 0x1000;
const INITIAL_ALLOC_PD_OFFSET: usize = 0x2000;
const INITIAL_ALLOC_OBJECT_POOL_PT_OFFSET: usize = 0x3000;
const INITIAL_ALLOC_KERNEL_PT_START_OFFSET: usize = 0x4000;

// Below should be used AFTER switching to new page table structure.
pub const OBJECT_POOL_START_VADDR: VAddr = VAddr::new(KERNEL_BASE + 0xe00000);
pub const OBJECT_POOL_SIZE: usize = 511;
pub const OBJECT_POOL_PT_VADDR: VAddr = VAddr::new(KERNEL_BASE + 0xfff000);

pub static APIC_PAGE: ExternMutex<APICPage> = unsafe { ExternMutex::new(None) };

// Variables
static INITIAL_PD: ExternMutex<PD> =
    unsafe { ExternMutex::new(Some(&init_pd as *const _)) };

pub static OBJECT_POOL_PT: ExternMutex<[PTEntry; OBJECT_POOL_SIZE]> =
    unsafe { ExternMutex::new(None) };
pub static KERNEL_PML4: ExternReadonlyObject<PML4> =
    unsafe { ExternReadonlyObject::new() };
pub static KERNEL_PDPT: ExternReadonlyObject<PDPT> =
    unsafe { ExternReadonlyObject::new() };
pub static KERNEL_PD: ExternReadonlyObject<PD> =
    unsafe { ExternReadonlyObject::new() };

fn kernel_stack_guard_page_vaddr() -> VAddr {
    VAddr::from((&kernel_stack_guard_page as *const _) as u64)
}

fn alloc_kernel_pml4(region: &mut MemoryRegion, alloc_base: PAddr) -> Unique<PML4> {
    use arch::paging::PML4Entry;

    let paddr = alloc_base + INITIAL_ALLOC_PML4_OFFSET;
    let vaddr = INITIAL_ALLOC_START_VADDR + INITIAL_ALLOC_PML4_OFFSET;

    log!("pml4, paddr: 0x{:x}, vaddr: 0x{:x}", paddr, vaddr);

    let mut pml4_unique = unsafe { Unique::new(vaddr.into(): usize as *mut PML4) };

    {
        let mut pml4 = unsafe { pml4_unique.get_mut() };
        *pml4 = [PML4Entry::empty(); 512];
    }

    region.move_up(paddr + BASE_PAGE_LENGTH);

    unsafe { KERNEL_PML4.bootstrap(*pml4_unique.deref(), paddr); }

    pml4_unique
}

fn alloc_kernel_pdpt(region: &mut MemoryRegion, pml4: &mut PML4, alloc_base: PAddr) -> Unique<PDPT> {
    use arch::paging::{PDPTEntry, PML4Entry, PML4_P, PML4_RW};

    let paddr = alloc_base + INITIAL_ALLOC_PDPT_OFFSET;
    let vaddr = INITIAL_ALLOC_START_VADDR + INITIAL_ALLOC_PDPT_OFFSET;

    log!("pdpt, paddr: 0x{:x}, vaddr: 0x{:x}", paddr, vaddr);

    let mut pdpt_unique = unsafe { Unique::new(vaddr.into(): usize as *mut PDPT) };

    {
        let mut pdpt = unsafe { pdpt_unique.get_mut() };
        *pdpt = [PDPTEntry::empty(); 512];
    }

    region.move_up(paddr + BASE_PAGE_LENGTH);

    pml4[pml4_index(VAddr::from(KERNEL_BASE))] = PML4Entry::new(paddr, PML4_P | PML4_RW);

    unsafe { KERNEL_PDPT.bootstrap(*pdpt_unique.deref(), paddr) }

    pdpt_unique
}

fn alloc_kernel_pd(region: &mut MemoryRegion, pdpt: &mut PDPT, alloc_base: PAddr) -> Unique<PD> {
    use arch::paging::{PDEntry, PDPTEntry, PDPT_P, PDPT_RW};

    let paddr = alloc_base + INITIAL_ALLOC_PD_OFFSET;
    let vaddr = INITIAL_ALLOC_START_VADDR + INITIAL_ALLOC_PD_OFFSET;

    log!("pd, paddr: 0x{:x}, vaddr: 0x{:x}", paddr, vaddr);

    let mut pd_unique = unsafe { Unique::new(vaddr.into(): usize as *mut PD) };

    {
        let mut pd = unsafe { pd_unique.get_mut() };
        *pd = [PDEntry::empty(); 512];
    }

    region.move_up(paddr + BASE_PAGE_LENGTH);

    pdpt[pdpt_index(VAddr::from(KERNEL_BASE))] = PDPTEntry::new(paddr, PDPT_P | PDPT_RW);

    unsafe { KERNEL_PD.bootstrap(*pd_unique.deref(), paddr); }

    pd_unique
}

fn alloc_object_pool_pt(region: &mut MemoryRegion, pd: &mut PD, alloc_base: PAddr) -> Unique<PT> {
    use arch::paging::{PTEntry, PDEntry, PD_P, PD_RW, PT_P, PT_RW};

    let paddr = alloc_base + INITIAL_ALLOC_OBJECT_POOL_PT_OFFSET;
    let vaddr = INITIAL_ALLOC_START_VADDR + INITIAL_ALLOC_OBJECT_POOL_PT_OFFSET;

    log!("object_pool_pt, paddr: 0x{:x}, vaddr: 0x{:x}", paddr, vaddr);

    let mut pt_unique = unsafe { Unique::new(vaddr.into(): usize as *mut PT) };

    {
        let mut pt = unsafe { pt_unique.get_mut() };
        *pt = [PTEntry::empty(); 512];

        let reverse_pt_index = pt_index(OBJECT_POOL_PT_VADDR);
        log!("reverse_pt_index: {}", reverse_pt_index);

        pt[reverse_pt_index] = PTEntry::new(paddr, PT_P | PT_RW);
    }

    region.move_up(paddr + BASE_PAGE_LENGTH);

    pd[pd_index(OBJECT_POOL_START_VADDR)] = PDEntry::new(paddr, PD_P | PD_RW);

    pt_unique
}

fn alloc_apic(pt: &mut PT, offset_size: usize, alloc_base: PAddr) {
    use arch::paging::{PT_P, PT_RW, PT_PWT, PT_PCD};
    use x86::shared::msr;
    use core::mem;

    let apic_msr = unsafe { msr::rdmsr(0x1B) };
    let apic_base = (apic_msr >> 12) * 0x1000;
    assert!(apic_msr & (0x1 << 11) == (0x1 << 11));
    log!("apic base address: 0x{:x}", apic_base);

    assert!(mem::size_of::<APICPage>() == 1024);

    let paddr = PAddr::from(apic_base);
    let vaddr = kernel_start_vaddr() + (offset_size * BASE_PAGE_LENGTH);

    log!("apic allocated at 0x{:x}", vaddr);

    pt[pt_index(vaddr)] = PTEntry::new(paddr, PT_P | PT_RW | PT_PWT | PT_PCD);

    unsafe { APIC_PAGE.bootstrap(vaddr.into(): usize as *const APICPage) };
}

fn alloc_kernel_page(pt: &mut PT, offset_size: usize, alloc_base: PAddr) {
    use arch::paging::{PT_P, PT_RW};

    let paddr = kernel_start_paddr() + (offset_size * BASE_PAGE_LENGTH);
    let vaddr = kernel_start_vaddr() + (offset_size * BASE_PAGE_LENGTH);

    log!("kernel page allocated at 0x{:x}", vaddr);

    pt[pt_index(vaddr)] = PTEntry::new(paddr, PT_P | PT_RW);
}

fn alloc_kernel_guard_page(pt: &mut PT, offset_size: usize, alloc_base: PAddr) {
    use arch::paging::{PT_P, PT_RW};

    let paddr = kernel_start_paddr() + (offset_size * BASE_PAGE_LENGTH);
    let vaddr = kernel_start_vaddr() + (offset_size * BASE_PAGE_LENGTH);

    log!("guard page allocated at 0x{:x}", vaddr);

    pt[pt_index(vaddr)] = PTEntry::empty();
}

fn alloc_kernel_pts(region: &mut MemoryRegion, pd: &mut PD, alloc_base: PAddr) {
    use arch::paging::{PDEntry, PD_P, PD_RW};

    let kernel_page_size = block_count(kernel_end_paddr().into(): usize -
                                       kernel_start_paddr().into(): usize, BASE_PAGE_LENGTH);
    let guard_page_index = (kernel_stack_guard_page_vaddr().into(): usize -
                            kernel_start_vaddr().into(): usize) / BASE_PAGE_LENGTH;

    log!("guard_page_index: {}", guard_page_index);

    for i in 0..(kernel_page_size + 1) {
        if i % 512 == 0 {
            pd[pd_index(kernel_start_vaddr() + i * BASE_PAGE_LENGTH)] = PDEntry::new(region.start_paddr(), PD_P | PD_RW);
            let npaddr = region.start_paddr() + BASE_PAGE_LENGTH;
            region.move_up(npaddr);
        }

        let pt_entry = pd[pd_index(kernel_start_vaddr() + i * BASE_PAGE_LENGTH)];

        let offset = pt_entry.get_address().into(): usize -
            alloc_base.into(): usize;
        let mut pt_unique = unsafe {
            Unique::new((INITIAL_ALLOC_START_VADDR + offset).into(): usize as *mut PT) };

        if i == guard_page_index {
            alloc_kernel_guard_page(unsafe { pt_unique.get_mut() }, i % 512, alloc_base);
        } else if i == kernel_page_size {
            alloc_apic(unsafe { pt_unique.get_mut() }, i % 512, alloc_base);
        } else {
            alloc_kernel_page(unsafe { pt_unique.get_mut() }, i % 512, alloc_base);
        }
    }
}

// This maps 2MB for allocation region
fn map_alloc_region(alloc_region: &mut MemoryRegion) -> PAddr {
    use arch::paging::{PD_P, PD_RW, PD_PS, PDEntry, flush_all};

    let map_alloc_start_vaddr = INITIAL_ALLOC_START_VADDR;
    let map_alloc_pd_index = pd_index(map_alloc_start_vaddr);
    let map_alloc_start_paddr = align_up(alloc_region.start_paddr(), LARGE_PAGE_LENGTH);

    let mut initial_pd = INITIAL_PD.lock();
    initial_pd[map_alloc_pd_index] = PDEntry::new(map_alloc_start_paddr, PD_P | PD_RW | PD_PS);

    // unsafe { super::paging::flush(map_alloc_start_vaddr); }
    unsafe { flush_all(); }

    alloc_region.move_up(map_alloc_start_paddr);

    map_alloc_start_paddr
}

pub fn init(mut alloc_region: &mut MemoryRegion) {
    use arch::paging::{switch_to};

    let kernel_page_size = block_count(kernel_end_paddr().into(): usize -
                                       kernel_start_paddr().into(): usize, BASE_PAGE_LENGTH);
    let alloc_size = 3 /* PML4, PDPT, and PD, and object pool PT */ +
        block_count(kernel_page_size, 512) /* Kernel page mapping PT */ ;

    assert!(alloc_size <= 512);
    assert!(alloc_size * BASE_PAGE_LENGTH < alloc_region.length());

    // Before allocation, we need to make sure PAddr + alloc_size is mapped.
    // This is done in the init_pd.

    let alloc_base_paddr = map_alloc_region(&mut alloc_region);

    log!("alloc_base_paddr: 0x{:x}", alloc_base_paddr);

    let mut pml4_unique = alloc_kernel_pml4(&mut alloc_region,
                                            alloc_base_paddr);
    let mut pdpt_unique = alloc_kernel_pdpt(&mut alloc_region,
                                            unsafe { pml4_unique.get_mut() },
                                            alloc_base_paddr);
    let mut pd_unique = alloc_kernel_pd(&mut alloc_region,
                                        unsafe { pdpt_unique.get_mut() },
                                        alloc_base_paddr);
    let mut object_pool_pt_unique = alloc_object_pool_pt(&mut alloc_region,
                                                         unsafe { pd_unique.get_mut() },
                                                         alloc_base_paddr);

    alloc_kernel_pts(&mut alloc_region, unsafe { pd_unique.get_mut() }, alloc_base_paddr);

    unsafe {
        INITIAL_PD.unbootstrap();
    }
    unsafe { switch_to(KERNEL_PML4.paddr()); }
    unsafe {
        OBJECT_POOL_PT.bootstrap(OBJECT_POOL_PT_VADDR.into(): usize as *mut _);
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 Amanieu d'Antras
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

use spinwait::SpinWait;
use std::cell::Cell;
use std::mem;
#[cfg(not(has_localkey_try_with))]
use std::panic;
use std::ptr;
use std::sync::atomic::{fence, AtomicUsize, Ordering};
use std::thread::LocalKey;
use thread_parker::ThreadParker;

struct ThreadData {
    parker: ThreadParker,

    // Linked list of threads in the queue. The queue is split into two parts:
    // the processed part and the unprocessed part. When new nodes are added to
    // the list, they only have the next pointer set, and queue_tail is null.
    //
    // Nodes are processed with the queue lock held, which consists of setting
    // the prev pointer for each node and setting the queue_tail pointer on the
    // first processed node of the list.
    //
    // This setup allows nodes to be added to the queue without a lock, while
    // still allowing O(1) removal of nodes from the processed part of the list.
    // The only cost is the O(n) processing, but this only needs to be done
    // once for each node, and therefore isn't too expensive.
    queue_tail: Cell<*const ThreadData>,
    prev: Cell<*const ThreadData>,
    next: Cell<*const ThreadData>,
}

impl ThreadData {
    fn new() -> ThreadData {
        ThreadData {
            parker: ThreadParker::new(),
            queue_tail: Cell::new(ptr::null()),
            prev: Cell::new(ptr::null()),
            next: Cell::new(ptr::null()),
        }
    }
}

// Returns a ThreadData structure for the current thread
unsafe fn get_thread_data(local: &mut Option<ThreadData>) -> &ThreadData {
    // Try to read from thread-local storage, but return None if the TLS has
    // already been destroyed.
    #[cfg(has_localkey_try_with)]
    fn try_get_tls(key: &'static LocalKey<ThreadData>) -> Option<*const ThreadData> {
        key.try_with(|x| x as *const ThreadData).ok()
    }
    #[cfg(not(has_localkey_try_with))]
    fn try_get_tls(key: &'static LocalKey<ThreadData>) -> Option<*const ThreadData> {
        panic::catch_unwind(|| key.with(|x| x as *const ThreadData)).ok()
    }

    // If ThreadData is expensive to construct, then we want to use a cached
    // version in thread-local storage if possible.
    if !cfg!(windows) && !cfg!(all(feature = "nightly", target_os = "linux")) {
        thread_local!(static THREAD_DATA: ThreadData = ThreadData::new());
        if let Some(tls) = try_get_tls(&THREAD_DATA) {
            return &*tls;
        }
    }

    // Otherwise just create a ThreadData on the stack
    *local = Some(ThreadData::new());
    local.as_ref().unwrap()
}

const LOCKED_BIT: usize = 1;
const QUEUE_LOCKED_BIT: usize = 2;
const QUEUE_MASK: usize = !3;

// Word-sized lock that is used to implement the parking_lot API. Since this
// can't use parking_lot, it instead manages its own queue of waiting threads.
pub struct WordLock {
    state: AtomicUsize,
}

impl WordLock {
    #[inline]
    pub fn new() -> WordLock {
        WordLock {
            state: AtomicUsize::new(0),
        }
    }

    #[inline]
    pub unsafe fn lock(&self) {
        if self
            .state
            .compare_exchange_weak(0, LOCKED_BIT, Ordering::Acquire, Ordering::Relaxed)
            .is_ok()
        {
            return;
        }
        self.lock_slow();
    }

    #[inline]
    pub unsafe fn unlock(&self) {
        let state = self.state.fetch_sub(LOCKED_BIT, Ordering::Release);
        if state.is_queue_locked() || state.queue_head().is_null() {
            return;
        }
        self.unlock_slow();
    }

    #[cold]
    #[inline(never)]
    unsafe fn lock_slow(&self) {
        let mut spinwait = SpinWait::new();
        let mut state = self.state.load(Ordering::Relaxed);
        loop {
            // Grab the lock if it isn't locked, even if there is a queue on it
            if !state.is_locked() {
                match self.state.compare_exchange_weak(
                    state,
                    state | LOCKED_BIT,
                    Ordering::Acquire,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => return,
                    Err(x) => state = x,
                }
                continue;
            }

            // If there is no queue, try spinning a few times
            if state.queue_head().is_null() && spinwait.spin() {
                state = self.state.load(Ordering::Relaxed);
                continue;
            }

            // Get our thread data and prepare it for parking
            let mut thread_data = None;
            let thread_data = get_thread_data(&mut thread_data);
            assert!(mem::align_of_val(thread_data) > !QUEUE_MASK);
            thread_data.parker.prepare_park();

            // Add our thread to the front of the queue
            let queue_head = state.queue_head();
            if queue_head.is_null() {
                thread_data.queue_tail.set(thread_data);
                thread_data.prev.set(ptr::null());
            } else {
                thread_data.queue_tail.set(ptr::null());
                thread_data.prev.set(ptr::null());
                thread_data.next.set(queue_head);
            }
            if let Err(x) = self.state.compare_exchange_weak(
                state,
                state.with_queue_head(thread_data),
                Ordering::Release,
                Ordering::Relaxed,
            ) {
                state = x;
                continue;
            }

            // Sleep until we are woken up by an unlock
            thread_data.parker.park();

            // Loop back and try locking again
            spinwait.reset();
            self.state.load(Ordering::Relaxed);
        }
    }

    #[cold]
    #[inline(never)]
    unsafe fn unlock_slow(&self) {
        let mut state = self.state.load(Ordering::Relaxed);
        loop {
            // We just unlocked the WordLock. Just check if there is a thread
            // to wake up. If the queue is locked then another thread is already
            // taking care of waking up a thread.
            if state.is_queue_locked() || state.queue_head().is_null() {
                return;
            }

            // Try to grab the queue lock
            match self.state.compare_exchange_weak(
                state,
                state | QUEUE_LOCKED_BIT,
                Ordering::Acquire,
                Ordering::Relaxed,
            ) {
                Ok(_) => break,
                Err(x) => state = x,
            }
        }

        // Now we have the queue lock and the queue is non-empty
        'outer: loop {
            // First, we need to fill in the prev pointers for any newly added
            // threads. We do this until we reach a node that we previously
            // processed, which has a non-null queue_tail pointer.
            let queue_head = state.queue_head();
            let mut queue_tail;
            let mut current = queue_head;
            loop {
                queue_tail = (*current).queue_tail.get();
                if !queue_tail.is_null() {
                    break;
                }
                let next = (*current).next.get();
                (*next).prev.set(current);
                current = next;
            }

            // Set queue_tail on the queue head to indicate that the whole list
            // has prev pointers set correctly.
            (*queue_head).queue_tail.set(queue_tail);

            // If the WordLock is locked, then there is no point waking up a
            // thread now. Instead we let the next unlocker take care of waking
            // up a thread.
            if state.is_locked() {
                match self.state.compare_exchange_weak(
                    state,
                    state & !QUEUE_LOCKED_BIT,
                    Ordering::Release,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => return,
                    Err(x) => state = x,
                }

                // Need an acquire fence before reading the new queue
                fence(Ordering::Acquire);
                continue;
            }

            // Remove the last thread from the queue and unlock the queue
            let new_tail = (*queue_tail).prev.get();
            if new_tail.is_null() {
                loop {
                    match self.state.compare_exchange_weak(
                        state,
                        state & LOCKED_BIT,
                        Ordering::Release,
                        Ordering::Relaxed,
                    ) {
                        Ok(_) => break,
                        Err(x) => state = x,
                    }

                    // If the compare_exchange failed because a new thread was
                    // added to the queue then we need to re-scan the queue to
                    // find the previous element.
                    if state.queue_head().is_null() {
                        continue;
                    } else {
                        // Need an acquire fence before reading the new queue
                        fence(Ordering::Acquire);
                        continue 'outer;
                    }
                }
            } else {
                (*queue_head).queue_tail.set(new_tail);
                self.state.fetch_and(!QUEUE_LOCKED_BIT, Ordering::Release);
            }

            // Finally, wake up the thread we removed from the queue. Note that
            // we don't need to worry about any races here since the thread is
            // guaranteed to be sleeping right now and we are the only one who
            // can wake it up.
            (*queue_tail).parker.unpark_lock().unpark();
            break;
        }
    }
}

trait LockState {
    fn is_locked(self) -> bool;
    fn is_queue_locked(self) -> bool;
    fn queue_head(self) -> *const ThreadData;
    fn with_queue_head(self, thread_data: *const ThreadData) -> Self;
}

impl LockState for usize {
    #[inline]
    fn is_locked(self) -> bool {
        self & LOCKED_BIT != 0
    }

    #[inline]
    fn is_queue_locked(self) -> bool {
        self & QUEUE_LOCKED_BIT != 0
    }

    #[inline]
    fn queue_head(self) -> *const ThreadData {
        (self & QUEUE_MASK) as *const ThreadData
    }

    #[inline]
    fn with_queue_head(self, thread_data: *const ThreadData) -> Self {
        (self & !QUEUE_MASK) | thread_data as *const _ as usize
    }
}
<commit_msg>Make get_thread_data safe<commit_after>// Copyright 2016 Amanieu d'Antras
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

use spinwait::SpinWait;
use std::cell::Cell;
use std::mem;
#[cfg(not(has_localkey_try_with))]
use std::panic;
use std::ptr;
use std::sync::atomic::{fence, AtomicUsize, Ordering};
use std::thread::LocalKey;
use thread_parker::ThreadParker;

struct ThreadData {
    parker: ThreadParker,

    // Linked list of threads in the queue. The queue is split into two parts:
    // the processed part and the unprocessed part. When new nodes are added to
    // the list, they only have the next pointer set, and queue_tail is null.
    //
    // Nodes are processed with the queue lock held, which consists of setting
    // the prev pointer for each node and setting the queue_tail pointer on the
    // first processed node of the list.
    //
    // This setup allows nodes to be added to the queue without a lock, while
    // still allowing O(1) removal of nodes from the processed part of the list.
    // The only cost is the O(n) processing, but this only needs to be done
    // once for each node, and therefore isn't too expensive.
    queue_tail: Cell<*const ThreadData>,
    prev: Cell<*const ThreadData>,
    next: Cell<*const ThreadData>,
}

impl ThreadData {
    fn new() -> ThreadData {
        ThreadData {
            parker: ThreadParker::new(),
            queue_tail: Cell::new(ptr::null()),
            prev: Cell::new(ptr::null()),
            next: Cell::new(ptr::null()),
        }
    }
}

// Returns a ThreadData structure for the current thread
fn get_thread_data(local: &mut Option<ThreadData>) -> &ThreadData {
    // Try to read from thread-local storage, but return None if the TLS has
    // already been destroyed.
    #[cfg(has_localkey_try_with)]
    fn try_get_tls(key: &'static LocalKey<ThreadData>) -> Option<*const ThreadData> {
        key.try_with(|x| x as *const ThreadData).ok()
    }
    #[cfg(not(has_localkey_try_with))]
    fn try_get_tls(key: &'static LocalKey<ThreadData>) -> Option<*const ThreadData> {
        panic::catch_unwind(|| key.with(|x| x as *const ThreadData)).ok()
    }

    // If ThreadData is expensive to construct, then we want to use a cached
    // version in thread-local storage if possible.
    if !cfg!(windows) && !cfg!(all(feature = "nightly", target_os = "linux")) {
        thread_local!(static THREAD_DATA: ThreadData = ThreadData::new());
        if let Some(tls) = try_get_tls(&THREAD_DATA) {
            return unsafe { &*tls };
        }
    }

    // Otherwise just create a ThreadData on the stack
    *local = Some(ThreadData::new());
    local.as_ref().unwrap()
}

const LOCKED_BIT: usize = 1;
const QUEUE_LOCKED_BIT: usize = 2;
const QUEUE_MASK: usize = !3;

// Word-sized lock that is used to implement the parking_lot API. Since this
// can't use parking_lot, it instead manages its own queue of waiting threads.
pub struct WordLock {
    state: AtomicUsize,
}

impl WordLock {
    #[inline]
    pub fn new() -> WordLock {
        WordLock {
            state: AtomicUsize::new(0),
        }
    }

    #[inline]
    pub unsafe fn lock(&self) {
        if self
            .state
            .compare_exchange_weak(0, LOCKED_BIT, Ordering::Acquire, Ordering::Relaxed)
            .is_ok()
        {
            return;
        }
        self.lock_slow();
    }

    #[inline]
    pub unsafe fn unlock(&self) {
        let state = self.state.fetch_sub(LOCKED_BIT, Ordering::Release);
        if state.is_queue_locked() || state.queue_head().is_null() {
            return;
        }
        self.unlock_slow();
    }

    #[cold]
    #[inline(never)]
    unsafe fn lock_slow(&self) {
        let mut spinwait = SpinWait::new();
        let mut state = self.state.load(Ordering::Relaxed);
        loop {
            // Grab the lock if it isn't locked, even if there is a queue on it
            if !state.is_locked() {
                match self.state.compare_exchange_weak(
                    state,
                    state | LOCKED_BIT,
                    Ordering::Acquire,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => return,
                    Err(x) => state = x,
                }
                continue;
            }

            // If there is no queue, try spinning a few times
            if state.queue_head().is_null() && spinwait.spin() {
                state = self.state.load(Ordering::Relaxed);
                continue;
            }

            // Get our thread data and prepare it for parking
            let mut thread_data = None;
            let thread_data = get_thread_data(&mut thread_data);
            assert!(mem::align_of_val(thread_data) > !QUEUE_MASK);
            thread_data.parker.prepare_park();

            // Add our thread to the front of the queue
            let queue_head = state.queue_head();
            if queue_head.is_null() {
                thread_data.queue_tail.set(thread_data);
                thread_data.prev.set(ptr::null());
            } else {
                thread_data.queue_tail.set(ptr::null());
                thread_data.prev.set(ptr::null());
                thread_data.next.set(queue_head);
            }
            if let Err(x) = self.state.compare_exchange_weak(
                state,
                state.with_queue_head(thread_data),
                Ordering::Release,
                Ordering::Relaxed,
            ) {
                state = x;
                continue;
            }

            // Sleep until we are woken up by an unlock
            thread_data.parker.park();

            // Loop back and try locking again
            spinwait.reset();
            self.state.load(Ordering::Relaxed);
        }
    }

    #[cold]
    #[inline(never)]
    unsafe fn unlock_slow(&self) {
        let mut state = self.state.load(Ordering::Relaxed);
        loop {
            // We just unlocked the WordLock. Just check if there is a thread
            // to wake up. If the queue is locked then another thread is already
            // taking care of waking up a thread.
            if state.is_queue_locked() || state.queue_head().is_null() {
                return;
            }

            // Try to grab the queue lock
            match self.state.compare_exchange_weak(
                state,
                state | QUEUE_LOCKED_BIT,
                Ordering::Acquire,
                Ordering::Relaxed,
            ) {
                Ok(_) => break,
                Err(x) => state = x,
            }
        }

        // Now we have the queue lock and the queue is non-empty
        'outer: loop {
            // First, we need to fill in the prev pointers for any newly added
            // threads. We do this until we reach a node that we previously
            // processed, which has a non-null queue_tail pointer.
            let queue_head = state.queue_head();
            let mut queue_tail;
            let mut current = queue_head;
            loop {
                queue_tail = (*current).queue_tail.get();
                if !queue_tail.is_null() {
                    break;
                }
                let next = (*current).next.get();
                (*next).prev.set(current);
                current = next;
            }

            // Set queue_tail on the queue head to indicate that the whole list
            // has prev pointers set correctly.
            (*queue_head).queue_tail.set(queue_tail);

            // If the WordLock is locked, then there is no point waking up a
            // thread now. Instead we let the next unlocker take care of waking
            // up a thread.
            if state.is_locked() {
                match self.state.compare_exchange_weak(
                    state,
                    state & !QUEUE_LOCKED_BIT,
                    Ordering::Release,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => return,
                    Err(x) => state = x,
                }

                // Need an acquire fence before reading the new queue
                fence(Ordering::Acquire);
                continue;
            }

            // Remove the last thread from the queue and unlock the queue
            let new_tail = (*queue_tail).prev.get();
            if new_tail.is_null() {
                loop {
                    match self.state.compare_exchange_weak(
                        state,
                        state & LOCKED_BIT,
                        Ordering::Release,
                        Ordering::Relaxed,
                    ) {
                        Ok(_) => break,
                        Err(x) => state = x,
                    }

                    // If the compare_exchange failed because a new thread was
                    // added to the queue then we need to re-scan the queue to
                    // find the previous element.
                    if state.queue_head().is_null() {
                        continue;
                    } else {
                        // Need an acquire fence before reading the new queue
                        fence(Ordering::Acquire);
                        continue 'outer;
                    }
                }
            } else {
                (*queue_head).queue_tail.set(new_tail);
                self.state.fetch_and(!QUEUE_LOCKED_BIT, Ordering::Release);
            }

            // Finally, wake up the thread we removed from the queue. Note that
            // we don't need to worry about any races here since the thread is
            // guaranteed to be sleeping right now and we are the only one who
            // can wake it up.
            (*queue_tail).parker.unpark_lock().unpark();
            break;
        }
    }
}

trait LockState {
    fn is_locked(self) -> bool;
    fn is_queue_locked(self) -> bool;
    fn queue_head(self) -> *const ThreadData;
    fn with_queue_head(self, thread_data: *const ThreadData) -> Self;
}

impl LockState for usize {
    #[inline]
    fn is_locked(self) -> bool {
        self & LOCKED_BIT != 0
    }

    #[inline]
    fn is_queue_locked(self) -> bool {
        self & QUEUE_LOCKED_BIT != 0
    }

    #[inline]
    fn queue_head(self) -> *const ThreadData {
        (self & QUEUE_MASK) as *const ThreadData
    }

    #[inline]
    fn with_queue_head(self, thread_data: *const ThreadData) -> Self {
        (self & !QUEUE_MASK) | thread_data as *const _ as usize
    }
}
<|endoftext|>


<commit_before>//! cfg defines conditional compiling options, `cfg` attibute parser and evaluator

mod cfg_expr;
mod dnf;
#[cfg(test)]
mod tests;

use std::fmt;

use rustc_hash::FxHashSet;
use tt::SmolStr;

pub use cfg_expr::{CfgAtom, CfgExpr};
pub use dnf::DnfExpr;

/// Configuration options used for conditional compilition on items with `cfg` attributes.
/// We have two kind of options in different namespaces: atomic options like `unix`, and
/// key-value options like `target_arch="x86"`.
///
/// Note that for key-value options, one key can have multiple values (but not none).
/// `feature` is an example. We have both `feature="foo"` and `feature="bar"` if features
/// `foo` and `bar` are both enabled. And here, we store key-value options as a set of tuple
/// of key and value in `key_values`.
///
/// See: https://doc.rust-lang.org/reference/conditional-compilation.html#set-configuration-options
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct CfgOptions {
    enabled: FxHashSet<CfgAtom>,
}

impl CfgOptions {
    pub fn check(&self, cfg: &CfgExpr) -> Option<bool> {
        cfg.fold(&|atom| self.enabled.contains(atom))
    }

    pub fn insert_atom(&mut self, key: SmolStr) {
        self.enabled.insert(CfgAtom::Flag(key));
    }

    pub fn insert_key_value(&mut self, key: SmolStr, value: SmolStr) {
        self.enabled.insert(CfgAtom::KeyValue { key, value });
    }

    pub fn append(&mut self, other: &CfgOptions) {
        for atom in &other.enabled {
            self.enabled.insert(atom.clone());
        }
    }

    pub fn apply_diff(&mut self, diff: CfgDiff) {
        for atom in diff.enable {
            self.enabled.insert(atom);
        }

        for atom in diff.disable {
            self.enabled.remove(&atom);
        }
    }
}

pub struct CfgDiff {
    // Invariants: No duplicates, no atom that's both in `enable` and `disable`.
    enable: Vec<CfgAtom>,
    disable: Vec<CfgAtom>,
}

impl CfgDiff {
    /// Returns the total number of atoms changed by this diff.
    pub fn len(&self) -> usize {
        self.enable.len() + self.disable.len()
    }
}

impl fmt::Display for CfgDiff {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if !self.enable.is_empty() {
            f.write_str("enable ")?;
            for (i, atom) in self.enable.iter().enumerate() {
                let sep = match i {
                    0 => "",
                    _ if i == self.enable.len() - 1 => " and ",
                    _ => ", ",
                };
                f.write_str(sep)?;

                write!(f, "{}", atom)?;
            }

            if !self.disable.is_empty() {
                f.write_str("; ")?;
            }
        }

        if !self.disable.is_empty() {
            f.write_str("disable ")?;
            for (i, atom) in self.disable.iter().enumerate() {
                let sep = match i {
                    0 => "",
                    _ if i == self.enable.len() - 1 => " and ",
                    _ => ", ",
                };
                f.write_str(sep)?;

                write!(f, "{}", atom)?;
            }
        }

        Ok(())
    }
}

pub struct InactiveReason {
    enabled: Vec<CfgAtom>,
    disabled: Vec<CfgAtom>,
}

impl fmt::Display for InactiveReason {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if !self.enabled.is_empty() {
            for (i, atom) in self.enabled.iter().enumerate() {
                let sep = match i {
                    0 => "",
                    _ if i == self.enabled.len() - 1 => " and ",
                    _ => ", ",
                };
                f.write_str(sep)?;

                write!(f, "{}", atom)?;
            }
            let is_are = if self.enabled.len() == 1 { "is" } else { "are" };
            write!(f, " {} enabled", is_are)?;

            if !self.disabled.is_empty() {
                f.write_str(" and ")?;
            }
        }

        if !self.disabled.is_empty() {
            for (i, atom) in self.disabled.iter().enumerate() {
                let sep = match i {
                    0 => "",
                    _ if i == self.disabled.len() - 1 => " and ",
                    _ => ", ",
                };
                f.write_str(sep)?;

                write!(f, "{}", atom)?;
            }
            let is_are = if self.disabled.len() == 1 { "is" } else { "are" };
            write!(f, " {} disabled", is_are)?;
        }

        Ok(())
    }
}
<commit_msg>Remove dead code<commit_after>//! cfg defines conditional compiling options, `cfg` attibute parser and evaluator

mod cfg_expr;
mod dnf;
#[cfg(test)]
mod tests;

use std::fmt;

use rustc_hash::FxHashSet;
use tt::SmolStr;

pub use cfg_expr::{CfgAtom, CfgExpr};
pub use dnf::DnfExpr;

/// Configuration options used for conditional compilition on items with `cfg` attributes.
/// We have two kind of options in different namespaces: atomic options like `unix`, and
/// key-value options like `target_arch="x86"`.
///
/// Note that for key-value options, one key can have multiple values (but not none).
/// `feature` is an example. We have both `feature="foo"` and `feature="bar"` if features
/// `foo` and `bar` are both enabled. And here, we store key-value options as a set of tuple
/// of key and value in `key_values`.
///
/// See: https://doc.rust-lang.org/reference/conditional-compilation.html#set-configuration-options
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct CfgOptions {
    enabled: FxHashSet<CfgAtom>,
}

impl CfgOptions {
    pub fn check(&self, cfg: &CfgExpr) -> Option<bool> {
        cfg.fold(&|atom| self.enabled.contains(atom))
    }

    pub fn insert_atom(&mut self, key: SmolStr) {
        self.enabled.insert(CfgAtom::Flag(key));
    }

    pub fn insert_key_value(&mut self, key: SmolStr, value: SmolStr) {
        self.enabled.insert(CfgAtom::KeyValue { key, value });
    }

    pub fn apply_diff(&mut self, diff: CfgDiff) {
        for atom in diff.enable {
            self.enabled.insert(atom);
        }

        for atom in diff.disable {
            self.enabled.remove(&atom);
        }
    }
}

pub struct CfgDiff {
    // Invariants: No duplicates, no atom that's both in `enable` and `disable`.
    enable: Vec<CfgAtom>,
    disable: Vec<CfgAtom>,
}

impl CfgDiff {
    /// Returns the total number of atoms changed by this diff.
    pub fn len(&self) -> usize {
        self.enable.len() + self.disable.len()
    }
}

impl fmt::Display for CfgDiff {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if !self.enable.is_empty() {
            f.write_str("enable ")?;
            for (i, atom) in self.enable.iter().enumerate() {
                let sep = match i {
                    0 => "",
                    _ if i == self.enable.len() - 1 => " and ",
                    _ => ", ",
                };
                f.write_str(sep)?;

                write!(f, "{}", atom)?;
            }

            if !self.disable.is_empty() {
                f.write_str("; ")?;
            }
        }

        if !self.disable.is_empty() {
            f.write_str("disable ")?;
            for (i, atom) in self.disable.iter().enumerate() {
                let sep = match i {
                    0 => "",
                    _ if i == self.enable.len() - 1 => " and ",
                    _ => ", ",
                };
                f.write_str(sep)?;

                write!(f, "{}", atom)?;
            }
        }

        Ok(())
    }
}

pub struct InactiveReason {
    enabled: Vec<CfgAtom>,
    disabled: Vec<CfgAtom>,
}

impl fmt::Display for InactiveReason {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if !self.enabled.is_empty() {
            for (i, atom) in self.enabled.iter().enumerate() {
                let sep = match i {
                    0 => "",
                    _ if i == self.enabled.len() - 1 => " and ",
                    _ => ", ",
                };
                f.write_str(sep)?;

                write!(f, "{}", atom)?;
            }
            let is_are = if self.enabled.len() == 1 { "is" } else { "are" };
            write!(f, " {} enabled", is_are)?;

            if !self.disabled.is_empty() {
                f.write_str(" and ")?;
            }
        }

        if !self.disabled.is_empty() {
            for (i, atom) in self.disabled.iter().enumerate() {
                let sep = match i {
                    0 => "",
                    _ if i == self.disabled.len() - 1 => " and ",
                    _ => ", ",
                };
                f.write_str(sep)?;

                write!(f, "{}", atom)?;
            }
            let is_are = if self.disabled.len() == 1 { "is" } else { "are" };
            write!(f, " {} disabled", is_are)?;
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>use std::process::Command;
use std::path::{Path, MAIN_SEPARATOR};
use std::fs;
use std::env;

use archive;
use manifest::{load as load_manifest, Manifest};

use errors::*;

pub struct BuildOptions<'a> {
    pub release: bool,
    pub features: Option<Vec<&'a str>>,
    pub all_features: bool,
    pub no_default_features: bool,
}

pub fn build(options: BuildOptions) -> Result<()> {
    let features = options.features.unwrap_or(vec![]);

    println!("Build -- release: {}, features: {}, all-features: {}, no-default-features: {}\n",
             options.release,
             match features.len() {
                 0 => "(no features)".to_string(),
                 _ => features.join(" "),
             },
             options.all_features,
             options.no_default_features);

    let cwd = env::current_dir()
        .chain_err(|| "Failed to find current directory")?;
    let manifest = load_manifest("vba-block.toml")
        .chain_err(|| "Failed to load manifest")?;

    build_targets(cwd, manifest).expect("Failed to build targets");

    // TEMP
    match run("build", &features) {
        Ok(stdout) => println!("{}", stdout),
        Err(stderr) => {
            println!("ERROR: {}", stderr);
            return Err(stderr);
        }
    }

    println!("Done.");
    Ok(())
}

fn build_targets<P: AsRef<Path>>(cwd: P, manifest: Manifest) -> Result<()> {
    let build_dir = cwd.as_ref().join("build");
    fs::create_dir_all(build_dir.clone())
        .chain_err(|| "Failed to create build folder")?;

    if let Some(targets) = manifest.targets {
        for target in targets {
            let name = target.name.unwrap_or(manifest.package.name.clone());
            let path = cwd.as_ref()
                .join(target
                          .path
                          .replace("/", MAIN_SEPARATOR.to_string().as_str()));
            let file = build_dir.join(format!("{}.{}", name, target.extension));

            archive::zip(path, file)
                .chain_err(|| "Failed to create target")?;
        }
    }

    Ok(())
}

fn run(name: &str, args: &Vec<&str>) -> Result<String> {
    let mut script_dir = env::current_exe()
        .chain_err(|| "Failed to find script directory")?;

    // TODO Check for development vs installed

    script_dir.pop(); // vba-blocks.exe
    script_dir.pop(); // debug/
    script_dir.pop(); // target/
    script_dir.push("scripts");
    let script_dir = script_dir
        .to_str()
        .expect("Failed to find script directory");

    let output = if cfg!(target_os = "windows") {
        let command = format!("cscript {}\\run.vbs {} {}",
                              script_dir,
                              name,
                              args.join(" "));
        println!("Command: {}", command);
        Command::new("cmd")
            .args(&["/C", &command])
            .output()
            .expect("Failed to execute script")
    } else {
        let command = format!("osascript {}/run.scpt {} {}",
                              script_dir,
                              name,
                              args.join(" "));
        println!("Command: {}", command);
        Command::new("sh")
            .args(&["-c", &command])
            .output()
            .expect("Failed to execute script")
    };

    if output.stderr.len() > 0 {
        bail!(String::from_utf8(output.stderr).unwrap());
    } else {
        Ok(String::from_utf8(output.stdout).unwrap())
    }
}
<commit_msg>Tweak error-handling<commit_after>use std::process::Command;
use std::path::{Path, MAIN_SEPARATOR};
use std::fs;
use std::env;

use archive;
use manifest::{load as load_manifest, Manifest};

use errors::*;

pub struct BuildOptions<'a> {
    pub release: bool,
    pub features: Option<Vec<&'a str>>,
    pub all_features: bool,
    pub no_default_features: bool,
}

pub fn build(options: BuildOptions) -> Result<()> {
    let features = options.features.unwrap_or(vec![]);

    println!("Build -- release: {}, features: {}, all-features: {}, no-default-features: {}\n",
             options.release,
             match features.len() {
                 0 => "(no features)".to_string(),
                 _ => features.join(" "),
             },
             options.all_features,
             options.no_default_features);

    let cwd = env::current_dir()
        .chain_err(|| "Failed to find current directory")?;
    let manifest = load_manifest("vba-block.toml")
        .chain_err(|| "Failed to load manifest")?;

    build_targets(cwd, manifest)?;

    // TEMP
    match run("build", &features) {
        Ok(stdout) => println!("{}", stdout),
        Err(stderr) => {
            println!("ERROR: {}", stderr);
            return Err(stderr);
        }
    }

    println!("Done.");
    Ok(())
}

fn build_targets<P: AsRef<Path>>(cwd: P, manifest: Manifest) -> Result<()> {
    let build_dir = cwd.as_ref().join("build");
    fs::create_dir_all(build_dir.clone())
        .chain_err(|| "Failed to create build folder")?;

    if let Some(targets) = manifest.targets {
        for target in targets {
            let name = target.name.unwrap_or(manifest.package.name.clone());
            let path = cwd.as_ref()
                .join(target
                          .path
                          .replace("/", MAIN_SEPARATOR.to_string().as_str()));
            let file = build_dir.join(format!("{}.{}", name, target.extension));

            archive::zip(path, file)
                .chain_err(|| "Failed to create target")?;
        }
    }

    Ok(())
}

fn run(name: &str, args: &Vec<&str>) -> Result<String> {
    let mut script_dir = env::current_exe()
        .chain_err(|| "Failed to find script directory")?;

    // TODO Check for development vs installed

    script_dir.pop(); // vba-blocks.exe
    script_dir.pop(); // debug/
    script_dir.pop(); // target/
    script_dir.push("scripts");
    let script_dir = script_dir
        .to_str()
        .ok_or("Failed to find script directory")?;

    let output = if cfg!(target_os = "windows") {
        let command = format!("cscript {}\\run.vbs {} {}",
                              script_dir,
                              name,
                              args.join(" "));
        println!("Command: {}", command);
        Command::new("cmd")
            .args(&["/C", &command])
            .output()
            .chain_err(|| "Failed to execute script")?
    } else {
        let command = format!("osascript {}/run.scpt {} {}",
                              script_dir,
                              name,
                              args.join(" "));
        println!("Command: {}", command);
        Command::new("sh")
            .args(&["-c", &command])
            .output()
            .chain_err(|| "Failed to execute script")?
    };

    if output.stderr.len() > 0 {
        bail!(String::from_utf8(output.stderr).unwrap());
    } else {
        Ok(String::from_utf8(output.stdout).unwrap())
    }
}
<|endoftext|>


<commit_before>fn main() {
    println!("Value in cents of Penny: {}", value_in_cents(Coin::Penny));
}

enum Coin {
    Penny,
    Nickel,
    Dime,
    Quarter
}

fn value_in_cents(coin: Coin) -> u32 {
    match coin {
        Coin::Penny => 1,
        Coin::Nickel => 5,
        Coin::Dime => 10,
        Coin::Quarter => 25,
    }
}<commit_msg>Example of using variants that also hold value with match<commit_after>fn main() {
    println!("Value in cents of Penny: {}", value_in_cents(Coin::Penny));
    println!(
        "Value in cents of Quarter: {}",
        value_in_cents(Coin::Quarter(UsState::California)));
}

#[derive(Debug)]
enum UsState {
    California,
    Georgia,
    Texas
}

enum Coin {
    Penny,
    Nickel,
    Dime,
    Quarter(UsState)
}

fn value_in_cents(coin: Coin) -> u32 {
    match coin {
        Coin::Penny => 1,
        Coin::Nickel => 5,
        Coin::Dime => 10,
        Coin::Quarter(state) => {
            println!("Quarter from state: {:?}", state);
            25
        },
    }
}<|endoftext|>


<commit_before>fn lastbutone<'a, T>(list: &'a ~[T]) -> Option<(&'a T, &'a T)> {
    if list.len() > 2 {
        Some((&list[list.len()-2], &list[list.len()-1]))
    } else {
        None
    }
}

fn main() {
    let list = ~['a', 'b', 'c', 'd', 'e'];
    println!("{:?}", lastbutone(&list));

    let list = ~['a'];
    println!("{:?}", lastbutone(&list));
}
<commit_msg>Redo Problem 02 with pattern matching<commit_after>fn lastbutone<'a, T>(list: &'a [T]) -> Option<(&'a T, &'a T)> {
    match list {
        [] | [_] => None,
        [.., ref x, ref y] => Some((x, y))
    }
}

fn main() {
    let list = ~['a', 'b', 'c', 'd', 'e'];
    println!("{:?}", lastbutone(list));

    let list = ~['a'];
    println!("{:?}", lastbutone(list));
}
<|endoftext|>


<commit_before>use std::env;
use std::error::Error;
use std::fs::File;
use std::io::prelude::*;
use std::path::Path;
use std::process;

fn isprint(ch: u8) -> bool {
    ch >= 32 && ch <= 126
}

fn attack_cipher(ciphertext: Vec<u8>, crib: Vec<u8>) {
    for (i, x) in ciphertext.iter().enumerate() {
        let mut old_key = x ^ crib[0];
        let mut j = 1;
        while j < crib.len() && (i+j) < ciphertext.len() {
            let key = ciphertext[i+j] ^ crib[j];
            if key != old_key { break; }
            old_key = key;
            j += 1;
            if j == crib.len() {
                println!("Found text at 0x{:x} (XOR key 0x{:02x})", i, key);
                print!("  preview: ");
                for ch in &ciphertext[i..i+50] {
                    let raw = ch ^ key;
                    if isprint(raw) {
                        print!("{}", raw as char);
                    } else {
                        print!(".");
                    }
                }
                println!("");
            }
        }
    }
}

fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() < 3 {
        println!("usage: roxor <file> <crib>");
        process::exit(1);
    }

    let path = Path::new(&args[1]);
    let crib: Vec<u8> = args[2].bytes().collect();

    if crib.len() < 2 {
        println!("error: crib too short");
        process::exit(1);
    }

    let mut file = match File::open(&path) {
        Err(why) => panic!("couldn't open {}: {}", path.display(), Error::description(&why)),
        Ok(file) => file,
    };

    let mut ciphertext: Vec<u8> = Vec::new();
    file.read_to_end(&mut ciphertext).unwrap();

    if ciphertext.len() < crib.len() {
        println!("error: ciphertext too short");
        process::exit(1);
    }

    attack_cipher(ciphertext, crib);
}
<commit_msg>Use map() for preview<commit_after>use std::env;
use std::error::Error;
use std::fs::File;
use std::io::prelude::*;
use std::path::Path;
use std::process;

fn attack_cipher(ciphertext: Vec<u8>, crib: Vec<u8>) {
    for (i, x) in ciphertext.iter().enumerate() {
        let mut old_key = x ^ crib[0];
        let mut j = 1;
        while j < crib.len() && (i+j) < ciphertext.len() {
            let key = ciphertext[i+j] ^ crib[j];
            if key != old_key { break; }
            old_key = key;
            j += 1;
            if j == crib.len() {
                println!("Found text at 0x{:x} (XOR key 0x{:02x})", i, key);
                print!("  preview: ");
                let preview: String = ciphertext[i..i+50]
                    .iter()
                    .map(|x| x ^ key)
                    .map(|x| if x >= 32 && x <= 126 {x as char} else {'.'})
                    .collect();
                println!("{}", preview);
            }
        }
    }
}

fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() < 3 {
        println!("usage: roxor <file> <crib>");
        process::exit(1);
    }

    let path = Path::new(&args[1]);
    let crib: Vec<u8> = args[2].bytes().collect();

    if crib.len() < 2 {
        println!("error: crib too short");
        process::exit(1);
    }

    let mut file = match File::open(&path) {
        Err(why) => panic!("couldn't open {}: {}", path.display(), Error::description(&why)),
        Ok(file) => file,
    };

    let mut ciphertext: Vec<u8> = Vec::new();
    file.read_to_end(&mut ciphertext).unwrap();

    if ciphertext.len() < crib.len() {
        println!("error: ciphertext too short");
        process::exit(1);
    }

    attack_cipher(ciphertext, crib);
}
<|endoftext|>


<commit_before>use std::env;
use std::fs::File;
use std::io::prelude::*;
use std::io::stderr;
use std::path::Path;
use std::process;

const PREVIEW_LEN: usize = 50;

#[derive(Debug, PartialEq)]
struct Match {
    offset: usize,
    key: u8,
    preview: String,
}

fn main() {
    let args: Vec<String> = env::args().skip(1).collect();
    if args.len() < 2 {
        writeln!(stderr(), "usage: roxor <file> <crib>").unwrap();
        process::exit(1);
    }

    let path = Path::new(&args[0]);
    let crib = args[1].as_bytes();

    if crib.len() < 2 {
        writeln!(stderr(), "error: crib too short").unwrap();
        process::exit(1);
    }

    let mut file = match File::open(&path) {
        Err(why) => panic!("couldn't open {}: {}", path.display(), why),
        Ok(file) => file,
    };

    let mut ciphertext: Vec<u8> = Vec::new();
    file.read_to_end(&mut ciphertext).unwrap();

    if ciphertext.len() < crib.len() {
        writeln!(stderr(), "error: ciphertext too short").unwrap();
        process::exit(1);
    }

    for m in attack_cipher(&ciphertext[..], crib) {
        println!("Found text at 0x{:x} (XOR key 0x{:02x})", m.offset, m.key);
        println!("  preview: {}", m.preview);
    }
}

fn attack_cipher(ciphertext: &[u8], crib: &[u8]) -> Vec<Match> {
    let mut matches: Vec<Match> = Vec::new();
    for (i, x) in ciphertext.iter().enumerate() {
        let mut old_key = x ^ crib[0];
        let mut j = 1;
        while j < crib.len() && (i + j) < ciphertext.len() {
            let key = ciphertext[i + j] ^ crib[j];
            if key != old_key {
                break;
            }
            old_key = key;
            j += 1;
            if j == crib.len() {
                let preview: String = ciphertext[i..]
                    .iter()
                    .take(PREVIEW_LEN)
                    .map(|x| x ^ key)
                    .map(|x| if x >= 32 && x <= 126 { x as char } else { '.' })
                    .collect();
                matches.push(Match {
                    offset: i,
                    key: key,
                    preview: preview,
                });
            }
        }
    }
    return matches;
}

#[cfg(test)]
mod tests {
    use super::{attack_cipher, Match};

    struct Test {
        ciphertext: &'static [u8],
        crib: &'static [u8],
        matches: Vec<Match>,
    }

    #[test]
    fn test_attack_cipher() {
        let tests = vec![
            Test {
                ciphertext: &[],
                crib: &[],
                matches: vec![],
            },
            Test {
                ciphertext: b"haystack",
                crib: b"needle",
                matches: vec![],
            },
            Test {
                ciphertext: b"needle in haystack",
                crib: b"needle",
                matches: vec![
                    Match { offset: 0, key: 0, preview: "needle in haystack".to_string() },
                ],
            },
            Test {
                ciphertext: b"a needle, another needle",
                crib: b"needle",
                matches: vec![
                    Match { offset: 2, key: 0, preview: "needle, another needle".to_string() },
                    Match { offset: 18, key: 0, preview: "needle".to_string() },
                ],
            },
            Test {
                ciphertext: b"\x23\x62\x2c\x27\x27\x26\x2e\x27\x6e\x62\x23\x2c\x2d\x36\x2a\x27\x30\x62\x2c\x27\x27\x26\x2e\x27",
                crib: b"needle",
                matches: vec![
                    Match { offset: 2, key: 0x42, preview: "needle, another needle".to_string() },
                    Match { offset: 18, key: 0x42, preview: "needle".to_string() },
                ],
            },
        ];

        for t in tests.iter() {
            assert_eq!(t.matches, attack_cipher(t.ciphertext, t.crib));
        }
    }
}
<commit_msg>Make Clippy happy<commit_after>use std::env;
use std::fs::File;
use std::io::prelude::*;
use std::path::Path;
use std::process;

const PREVIEW_LEN: usize = 50;

#[derive(Debug, PartialEq)]
struct Match {
    offset: usize,
    key: u8,
    preview: String,
}

fn main() {
    let args: Vec<String> = env::args().skip(1).collect();
    if args.len() < 2 {
        eprintln!("usage: roxor <file> <crib>");
        process::exit(1);
    }

    let path = Path::new(&args[0]);
    let crib = args[1].as_bytes();

    if crib.len() < 2 {
        eprintln!("error: crib too short");
        process::exit(1);
    }

    let mut file = match File::open(&path) {
        Err(why) => panic!("couldn't open {}: {}", path.display(), why),
        Ok(file) => file,
    };

    let mut ciphertext: Vec<u8> = Vec::new();
    file.read_to_end(&mut ciphertext).unwrap();

    if ciphertext.len() < crib.len() {
        eprintln!("error: ciphertext too short");
        process::exit(1);
    }

    for m in attack_cipher(&ciphertext[..], crib) {
        println!("Found text at 0x{:x} (XOR key 0x{:02x})", m.offset, m.key);
        println!("  preview: {}", m.preview);
    }
}

fn attack_cipher(ciphertext: &[u8], crib: &[u8]) -> Vec<Match> {
    let mut matches: Vec<Match> = Vec::new();
    for (i, x) in ciphertext.iter().enumerate() {
        let mut old_key = x ^ crib[0];
        let mut j = 1;
        while j < crib.len() && (i + j) < ciphertext.len() {
            let key = ciphertext[i + j] ^ crib[j];
            if key != old_key {
                break;
            }
            old_key = key;
            j += 1;
            if j == crib.len() {
                let preview: String = ciphertext[i..]
                    .iter()
                    .take(PREVIEW_LEN)
                    .map(|x| x ^ key)
                    .map(|x| if x >= 32 && x <= 126 { x as char } else { '.' })
                    .collect();
                matches.push(Match {
                    offset: i,
                    key,
                    preview,
                });
            }
        }
    }
    matches
}

#[cfg(test)]
mod tests {
    use super::{attack_cipher, Match};

    struct Test {
        ciphertext: &'static [u8],
        crib: &'static [u8],
        matches: Vec<Match>,
    }

    #[test]
    fn test_attack_cipher() {
        let tests = vec![
            Test {
                ciphertext: &[],
                crib: &[],
                matches: vec![],
            },
            Test {
                ciphertext: b"haystack",
                crib: b"needle",
                matches: vec![],
            },
            Test {
                ciphertext: b"needle in haystack",
                crib: b"needle",
                matches: vec![
                    Match { offset: 0, key: 0, preview: "needle in haystack".to_string() },
                ],
            },
            Test {
                ciphertext: b"a needle, another needle",
                crib: b"needle",
                matches: vec![
                    Match { offset: 2, key: 0, preview: "needle, another needle".to_string() },
                    Match { offset: 18, key: 0, preview: "needle".to_string() },
                ],
            },
            Test {
                ciphertext: b"\x23\x62\x2c\x27\x27\x26\x2e\x27\x6e\x62\x23\x2c\x2d\x36\x2a\x27\x30\x62\x2c\x27\x27\x26\x2e\x27",
                crib: b"needle",
                matches: vec![
                    Match { offset: 2, key: 0x42, preview: "needle, another needle".to_string() },
                    Match { offset: 18, key: 0x42, preview: "needle".to_string() },
                ],
            },
        ];

        for t in tests.iter() {
            assert_eq!(t.matches, attack_cipher(t.ciphertext, t.crib));
        }
    }
}
<|endoftext|>


<commit_before>use std::{
    collections::VecDeque,
    fmt, io,
    sync::{
        atomic::{AtomicU64, AtomicUsize, Ordering},
        Arc, Mutex,
    },
    time::{Duration, Instant},
};

use crate::{
    config::ServerConfig,
    context::SharedContext,
    relay::tcprelay::EncryptedWrite,
    relay::{loadbalancing::server::LoadBalancer, socks5::Address, tcprelay::client::ServerClient},
};

use futures::{Future, Stream};
use log::{debug, error, info};
use tokio::{
    self,
    io::read_exact,
    timer::{Interval, Timeout},
};

struct Server {
    config: Arc<ServerConfig>,
    elapsed: AtomicU64,
}

impl Server {
    fn new(config: ServerConfig) -> Server {
        Server {
            config: Arc::new(config),
            elapsed: AtomicU64::new(0),
        }
    }

    fn delay(&self) -> u64 {
        self.elapsed.load(Ordering::Acquire)
    }
}

const MAX_LATENCY_QUEUE_SIZE: usize = 17;

struct ServerLatencyInner {
    latency_queue: VecDeque<u64>,
}

impl ServerLatencyInner {
    fn new() -> ServerLatencyInner {
        ServerLatencyInner {
            latency_queue: VecDeque::with_capacity(MAX_LATENCY_QUEUE_SIZE),
        }
    }

    fn push(&mut self, lat: u64) -> u64 {
        self.latency_queue.push_back(lat);
        if self.latency_queue.len() > MAX_LATENCY_QUEUE_SIZE {
            self.latency_queue.pop_front();
        }

        self.representation()
    }

    fn representation(&self) -> u64 {
        if self.latency_queue.is_empty() {
            return 0;
        }

        let sum: u64 = self.latency_queue.iter().sum();
        let avg = sum as f64 / self.latency_queue.len() as f64;
        let s_sum: u64 = self.latency_queue.iter().map(|n| *n * *n).sum();
        let dev = (s_sum as f64 / self.latency_queue.len() as f64 - avg * avg).sqrt();

        let mut v = self
            .latency_queue
            .iter()
            .cloned()
            .filter(|n| *n as f64 >= (avg - dev) && *n as f64 <= (avg + dev))
            .collect::<Vec<u64>>();
        v.sort();

        debug!(
            "Latency queue {:?}, avg {}, dev {}, sorted {:?}",
            self.latency_queue, avg, dev, v
        );

        if v.is_empty() {
            return 0;
        }

        let mid = v.len() / 2;
        if (v.len() & 1) == 0 {
            (v[mid - 1] + v[mid]) / 2
        } else {
            v[mid]
        }
    }
}

#[derive(Clone)]
struct ServerLatency {
    inner: Arc<Mutex<ServerLatencyInner>>,
}

impl ServerLatency {
    fn new() -> ServerLatency {
        ServerLatency {
            inner: Arc::new(Mutex::new(ServerLatencyInner::new())),
        }
    }

    fn push(&self, lat: u64) -> u64 {
        let mut inner = self.inner.lock().unwrap();
        inner.push(lat)
    }
}

impl fmt::Debug for ServerLatency {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let inner = self.inner.lock().unwrap();
        write!(f, "ServerLatency {{ latency: {:?} }}", inner.latency_queue)
    }
}

const DEFAULT_CHECK_INTERVAL_SEC: u64 = 10;
const DEFAULT_CHECK_TIMEOUT_SEC: u64 = 5;

struct Inner {
    servers: Vec<Arc<Server>>,
    best_idx: AtomicUsize,
}

impl Inner {
    fn new(context: SharedContext) -> Inner {
        let config = context.config();

        if config.server.is_empty() {
            panic!("no servers");
        }

        let check_required = config.server.len() > 1;

        let mut servers = Vec::new();
        for svr in &config.server {
            let sc = Arc::new(Server::new(svr.clone()));
            servers.push(sc.clone());

            if !check_required {
                continue;
            }

            let context = context.clone();
            let latency = ServerLatency::new();

            tokio::spawn(
                // Check every 10 seconds
                Interval::new(
                    Instant::now() + Duration::from_secs(1),
                    Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC),
                )
                .for_each(move |_| {
                    let sc = sc.clone();
                    let lat = latency.clone();

                    Inner::check_delay(sc.clone(), context.clone()).then(move |res| {
                        let md = match res {
                            Ok(d) => lat.push(d),
                            Err(..) => lat.push(DEFAULT_CHECK_TIMEOUT_SEC * 2 * 1000), // Penalty
                        };
                        debug!(
                            "updated remote server {} (median: {} ms) {:?}",
                            sc.config.addr(),
                            md,
                            lat
                        );
                        sc.elapsed.store(md, Ordering::Release);

                        Ok(())
                    })
                })
                .map_err(|_| ()),
            );
        }

        Inner {
            servers,
            best_idx: AtomicUsize::new(0),
        }
    }

    fn check_delay(sc: Arc<Server>, context: SharedContext) -> impl Future<Item = u64, Error = io::Error> {
        static GET_BODY: &[u8] = b"GET / HTTP/1.1\r\nHost: www.example.com\r\nConnection: close\r\nAccept: */*\r\n\r\n";

        let addr = Address::DomainNameAddress("www.example.com".to_owned(), 80);

        let start = Instant::now();

        // Send HTTP GET and read the first byte
        let fut = ServerClient::connect(context, addr, sc.config.clone()).and_then(|ServerClient { r, w }| {
            w.write_all(GET_BODY)
                .and_then(|(mut w, _)| w.flush())
                .and_then(move |_| read_exact(r, [0u8, 1]))
        });

        Timeout::new(fut, Duration::from_secs(DEFAULT_CHECK_TIMEOUT_SEC)).then(move |res| {
            let elapsed = Instant::now() - start;
            let elapsed = elapsed.as_secs() * 1000 + u64::from(elapsed.subsec_millis()); // Converted to ms
            match res {
                Ok(..) => {
                    // Got the result ... record its time
                    debug!("checked remote server {} latency with {} ms", sc.config.addr(), elapsed);
                    Ok(elapsed)
                }
                Err(err) => {
                    match err.into_inner() {
                        Some(err) => {
                            error!("failed to check server {}, error: {}", sc.config.addr(), err);

                            // NOTE: connection / handshake error, server is down
                            Err(err)
                        }
                        None => {
                            // Timeout
                            debug!(
                                "checked remote server {} latency timeout, elapsed {} ms",
                                sc.config.addr(),
                                elapsed
                            );

                            // NOTE: timeout is still available, but server is too slow
                            Ok(elapsed)
                        }
                    }
                }
            }
        })
    }

    fn best_idx(&self) -> usize {
        self.best_idx.load(Ordering::Acquire)
    }

    fn set_best_idx(&self, idx: usize) {
        self.best_idx.store(idx, Ordering::Release)
    }

    fn best_server(&self) -> &Arc<Server> {
        &self.servers[self.best_idx()]
    }
}

#[derive(Clone)]
pub struct PingBalancer {
    inner: Arc<Inner>,
}

impl PingBalancer {
    pub fn new(context: SharedContext) -> PingBalancer {
        let inner = Arc::new(Inner::new(context));

        let inner_cloned = inner.clone();
        tokio::spawn(
            Interval::new(
                Instant::now() + Duration::from_secs(2),
                Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC),
            )
            .for_each(move |_| {
                let inner = inner_cloned.clone();

                if inner.servers.is_empty() {
                    panic!("No server");
                }

                // Choose the best one
                let mut svr_idx = 0;

                for (idx, svr) in inner.servers.iter().enumerate() {
                    let choosen_svr = &inner.servers[svr_idx];
                    if choosen_svr.delay() == 0 {
                        continue;
                    }

                    if svr.delay() < choosen_svr.delay() {
                        svr_idx = idx;
                    }
                }

                let choosen_svr = &inner.servers[svr_idx];
                debug!(
                    "chosen the best server {} (delay: {} ms)",
                    choosen_svr.config.addr(),
                    choosen_svr.delay()
                );

                if inner.best_idx() != svr_idx {
                    info!(
                        "switched server from {} (latency: {} ms) to {} (latency: {} ms)",
                        inner.best_server().config.addr(),
                        inner.best_server().delay(),
                        choosen_svr.config.addr(),
                        choosen_svr.delay(),
                    );
                }

                inner.set_best_idx(svr_idx);

                Ok(())
            })
            .map_err(|_| ()),
        );

        PingBalancer { inner }
    }
}

impl LoadBalancer for PingBalancer {
    fn pick_server(&mut self) -> Arc<ServerConfig> {
        if self.inner.servers.is_empty() {
            panic!("no server in configuration");
        }
        self.inner.best_server().config.clone()
    }

    fn total(&self) -> usize {
        self.inner.servers.len()
    }
}
<commit_msg>Skip server if deviation is too large<commit_after>use std::{
    collections::VecDeque,
    fmt, io,
    sync::{
        atomic::{AtomicU64, AtomicUsize, Ordering},
        Arc, Mutex,
    },
    time::{Duration, Instant},
};

use crate::{
    config::ServerConfig,
    context::SharedContext,
    relay::tcprelay::EncryptedWrite,
    relay::{loadbalancing::server::LoadBalancer, socks5::Address, tcprelay::client::ServerClient},
};

use futures::{Future, Stream};
use log::{debug, error, info};
use tokio::{
    self,
    io::read_exact,
    timer::{Interval, Timeout},
};

struct Server {
    config: Arc<ServerConfig>,
    score: AtomicU64,
}

impl Server {
    fn new(config: ServerConfig) -> Server {
        Server {
            config: Arc::new(config),
            score: AtomicU64::new(0),
        }
    }

    fn score(&self) -> u64 {
        self.score.load(Ordering::Acquire)
    }
}

const MAX_LATENCY_QUEUE_SIZE: usize = 37;

struct ServerLatencyInner {
    latency_queue: VecDeque<u64>,
}

impl ServerLatencyInner {
    fn new() -> ServerLatencyInner {
        ServerLatencyInner {
            latency_queue: VecDeque::with_capacity(MAX_LATENCY_QUEUE_SIZE),
        }
    }

    fn push(&mut self, lat: u64) -> u64 {
        self.latency_queue.push_back(lat);
        if self.latency_queue.len() > MAX_LATENCY_QUEUE_SIZE {
            self.latency_queue.pop_front();
        }

        self.score()
    }

    fn score(&self) -> u64 {
        if self.latency_queue.is_empty() {
            return u64::max_value();
        }

        let sum: u64 = self.latency_queue.iter().sum();
        let avg = sum as f64 / self.latency_queue.len() as f64;
        let s_sum: u64 = self.latency_queue.iter().map(|n| *n * *n).sum();
        let dev = (s_sum as f64 / self.latency_queue.len() as f64 - avg * avg).sqrt();

        let mut v = self
            .latency_queue
            .iter()
            .cloned()
            .filter(|n| *n as f64 >= (avg - 3.0 * dev) && *n as f64 <= (avg + 3.0 * dev))
            .collect::<Vec<u64>>();
        v.sort();

        if v.is_empty() {
            return u64::max_value();
        }

        let mid = v.len() / 2;
        let median = if (v.len() & 1) == 0 {
            (v[mid - 1] + v[mid]) / 2
        } else {
            v[mid]
        };

        debug!(
            "latency queue {:?}, avg {}, dev {}, sorted {:?}, median {}",
            self.latency_queue, avg, dev, v, median
        );

        (((median as f64 * 2.0) + (dev * 1.0)) / 2.0) as u64
    }
}

#[derive(Clone)]
struct ServerLatency {
    inner: Arc<Mutex<ServerLatencyInner>>,
}

impl ServerLatency {
    fn new() -> ServerLatency {
        ServerLatency {
            inner: Arc::new(Mutex::new(ServerLatencyInner::new())),
        }
    }

    fn push(&self, lat: u64) -> u64 {
        let mut inner = self.inner.lock().unwrap();
        inner.push(lat)
    }
}

impl fmt::Debug for ServerLatency {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let inner = self.inner.lock().unwrap();
        write!(f, "ServerLatency {{ latency: {:?} }}", inner.latency_queue)
    }
}

const DEFAULT_CHECK_INTERVAL_SEC: u64 = 6;
const DEFAULT_CHECK_TIMEOUT_SEC: u64 = 5;

struct Inner {
    servers: Vec<Arc<Server>>,
    best_idx: AtomicUsize,
}

impl Inner {
    fn new(context: SharedContext) -> Inner {
        let config = context.config();

        if config.server.is_empty() {
            panic!("no servers");
        }

        let check_required = config.server.len() > 1;

        let mut servers = Vec::new();
        for svr in &config.server {
            let sc = Arc::new(Server::new(svr.clone()));
            servers.push(sc.clone());

            if !check_required {
                continue;
            }

            let context = context.clone();
            let latency = ServerLatency::new();

            tokio::spawn(
                // Check every DEFAULT_CHECK_INTERVAL_SEC seconds
                Interval::new(
                    Instant::now() + Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC),
                    Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC),
                )
                .for_each(move |_| {
                    let sc = sc.clone();
                    let lat = latency.clone();

                    Inner::check_delay(sc.clone(), context.clone()).then(move |res| {
                        let score = match res {
                            Ok(d) => lat.push(d),
                            Err(..) => lat.push(DEFAULT_CHECK_TIMEOUT_SEC * 2 * 1000), // Penalty
                        };
                        debug!("updated remote server {} (score: {})", sc.config.addr(), score);
                        sc.score.store(score, Ordering::Release);

                        Ok(())
                    })
                })
                .map_err(|_| ()),
            );
        }

        Inner {
            servers,
            best_idx: AtomicUsize::new(0),
        }
    }

    fn check_delay(sc: Arc<Server>, context: SharedContext) -> impl Future<Item = u64, Error = io::Error> {
        static GET_BODY: &[u8] = b"GET / HTTP/1.1\r\nHost: www.example.com\r\nConnection: close\r\nAccept: */*\r\n\r\n";

        let addr = Address::DomainNameAddress("www.example.com".to_owned(), 80);

        let start = Instant::now();

        // Send HTTP GET and read the first byte
        let fut = ServerClient::connect(context, addr, sc.config.clone()).and_then(|ServerClient { r, w }| {
            w.write_all(GET_BODY)
                .and_then(|(mut w, _)| w.flush())
                .and_then(move |_| read_exact(r, [0u8, 1]))
        });

        Timeout::new(fut, Duration::from_secs(DEFAULT_CHECK_TIMEOUT_SEC)).then(move |res| {
            let elapsed = Instant::now() - start;
            let elapsed = elapsed.as_secs() * 1000 + u64::from(elapsed.subsec_millis()); // Converted to ms
            match res {
                Ok(..) => {
                    // Got the result ... record its time
                    debug!("checked remote server {} latency with {} ms", sc.config.addr(), elapsed);
                    Ok(elapsed)
                }
                Err(err) => {
                    match err.into_inner() {
                        Some(err) => {
                            error!("failed to check server {}, error: {}", sc.config.addr(), err);

                            // NOTE: connection / handshake error, server is down
                            Err(err)
                        }
                        None => {
                            // Timeout
                            debug!(
                                "checked remote server {} latency timeout, elapsed {} ms",
                                sc.config.addr(),
                                elapsed
                            );

                            // NOTE: timeout is still available, but server is too slow
                            Ok(elapsed)
                        }
                    }
                }
            }
        })
    }

    fn best_idx(&self) -> usize {
        self.best_idx.load(Ordering::Acquire)
    }

    fn set_best_idx(&self, idx: usize) {
        self.best_idx.store(idx, Ordering::Release)
    }

    fn best_server(&self) -> &Arc<Server> {
        &self.servers[self.best_idx()]
    }
}

#[derive(Clone)]
pub struct PingBalancer {
    inner: Arc<Inner>,
}

impl PingBalancer {
    pub fn new(context: SharedContext) -> PingBalancer {
        let inner = Arc::new(Inner::new(context));

        let inner_cloned = inner.clone();
        tokio::spawn(
            Interval::new(
                Instant::now() + Duration::from_secs(2),
                Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC),
            )
            .for_each(move |_| {
                let inner = inner_cloned.clone();

                if inner.servers.is_empty() {
                    panic!("No server");
                }

                // Choose the best one
                let mut svr_idx = 0;

                for (idx, svr) in inner.servers.iter().enumerate() {
                    let choosen_svr = &inner.servers[svr_idx];
                    if svr.score() < choosen_svr.score() {
                        svr_idx = idx;
                    }
                }

                let choosen_svr = &inner.servers[svr_idx];
                debug!(
                    "chosen the best server {} (score: {})",
                    choosen_svr.config.addr(),
                    choosen_svr.score()
                );

                if inner.best_idx() != svr_idx {
                    info!(
                        "switched server from {} (score: {}) to {} (score: {})",
                        inner.best_server().config.addr(),
                        inner.best_server().score(),
                        choosen_svr.config.addr(),
                        choosen_svr.score(),
                    );
                }

                inner.set_best_idx(svr_idx);

                Ok(())
            })
            .map_err(|_| ()),
        );

        PingBalancer { inner }
    }
}

impl LoadBalancer for PingBalancer {
    fn pick_server(&mut self) -> Arc<ServerConfig> {
        if self.inner.servers.is_empty() {
            panic!("no server in configuration");
        }
        self.inner.best_server().config.clone()
    }

    fn total(&self) -> usize {
        self.inner.servers.len()
    }
}
<|endoftext|>