<commit_before>// Copyright 2020 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
use acpi_tables::{facs::FACS, rsdp::RSDP, sdt::SDT};
use data_model::DataInit;
use vm_memory::{GuestAddress, GuestMemory};

pub struct ACPIDevResource {
    pub amls: Vec<u8>,
    pub pm_iobase: u64,
    /// Additional system descriptor tables.
    pub sdts: Vec<SDT>,
}

#[repr(C)]
#[derive(Clone, Copy, Default)]
struct LocalAPIC {
    _type: u8,
    _length: u8,
    _processor_id: u8,
    _apic_id: u8,
    _flags: u32,
}

// Safe as LocalAPIC structure only contains raw data
unsafe impl DataInit for LocalAPIC {}

#[repr(C)]
#[derive(Clone, Copy, Default)]
struct IOAPIC {
    _type: u8,
    _length: u8,
    _ioapic_id: u8,
    _reserved: u8,
    _apic_address: u32,
    _gsi_base: u32,
}

// Safe as IOAPIC structure only contains raw data
unsafe impl DataInit for IOAPIC {}

const OEM_REVISION: u32 = 1;
//DSDT
const DSDT_REVISION: u8 = 6;
// FADT
const FADT_LEN: u32 = 276;
const FADT_REVISION: u8 = 6;
const FADT_MINOR_REVISION: u8 = 3;
// FADT flags
const FADT_POWER_BUTTON: u32 = 1 << 4;
const FADT_SLEEP_BUTTON: u32 = 1 << 5;
// FADT fields offset
const FADT_FIELD_FACS_ADDR32: usize = 36;
const FADT_FIELD_DSDT_ADDR32: usize = 40;
const FADT_FIELD_SCI_INTERRUPT: usize = 46;
const FADT_FIELD_PM1A_EVENT_BLK_ADDR: usize = 56;
const FADT_FIELD_PM1A_CONTROL_BLK_ADDR: usize = 64;
const FADT_FIELD_PM1A_EVENT_BLK_LEN: usize = 88;
const FADT_FIELD_PM1A_CONTROL_BLK_LEN: usize = 89;
const FADT_FIELD_FLAGS: usize = 112;
const FADT_FIELD_MINOR_REVISION: usize = 131;
const FADT_FIELD_FACS_ADDR: usize = 132;
const FADT_FIELD_DSDT_ADDR: usize = 140;
const FADT_FIELD_HYPERVISOR_ID: usize = 268;
// MADT
const MADT_LEN: u32 = 44;
const MADT_REVISION: u8 = 5;
// MADT fields offset
const MADT_FIELD_LAPIC_ADDR: usize = 36;
// MADT types
const MADT_TYPE_LOCAL_APIC: u8 = 0;
const MADT_TYPE_IO_APIC: u8 = 1;
// MADT flags
const MADT_ENABLED: u32 = 1;
// XSDT
const XSDT_REVISION: u8 = 1;

fn create_dsdt_table(amls: Vec<u8>) -> SDT {
    let mut dsdt = SDT::new(
        *b"DSDT",
        acpi_tables::HEADER_LEN,
        DSDT_REVISION,
        *b"CROSVM",
        *b"CROSVMDT",
        OEM_REVISION,
    );

    if !amls.is_empty() {
        dsdt.append_slice(amls.as_slice());
    }

    dsdt
}

fn create_facp_table(sci_irq: u16, pm_iobase: u32) -> SDT {
    let mut facp = SDT::new(
        *b"FACP",
        FADT_LEN,
        FADT_REVISION,
        *b"CROSVM",
        *b"CROSVMDT",
        OEM_REVISION,
    );

    let fadt_flags: u32 = FADT_POWER_BUTTON | FADT_SLEEP_BUTTON; // mask POWER and SLEEP BUTTON
    facp.write(FADT_FIELD_FLAGS, fadt_flags);

    // SCI Interrupt
    facp.write(FADT_FIELD_SCI_INTERRUPT, sci_irq);

    // PM1A Event Block Address
    facp.write(FADT_FIELD_PM1A_EVENT_BLK_ADDR, pm_iobase);

    // PM1A Control Block Address
    facp.write(
        FADT_FIELD_PM1A_CONTROL_BLK_ADDR,
        pm_iobase + devices::acpi::ACPIPM_RESOURCE_EVENTBLK_LEN as u32,
    );

    // PM1 Event Block Length
    facp.write(
        FADT_FIELD_PM1A_EVENT_BLK_LEN,
        devices::acpi::ACPIPM_RESOURCE_EVENTBLK_LEN as u8,
    );

    // PM1 Control Block Length
    facp.write(
        FADT_FIELD_PM1A_CONTROL_BLK_LEN,
        devices::acpi::ACPIPM_RESOURCE_CONTROLBLK_LEN as u8,
    );

    facp.write(FADT_FIELD_MINOR_REVISION, FADT_MINOR_REVISION); // FADT minor version
    facp.write(FADT_FIELD_HYPERVISOR_ID, *b"CROSVM"); // Hypervisor Vendor Identity

    facp
}

fn next_offset(offset: GuestAddress, len: u64) -> Option<GuestAddress> {
    // Enforce 64-byte allocation alignment.
    match len % 64 {
        0 => offset.checked_add(len),
        x => offset.checked_add(len.checked_add(64 - x)?),
    }
}

/// Create ACPI tables and return the RSDP.
/// The basic tables DSDT/FACP/MADT/XSDT are constructed in this function.
/// # Arguments
///
/// * `guest_mem` - The guest memory where the tables will be stored.
/// * `num_cpus` - Used to construct the MADT.
/// * `sci_irq` - Used to fill the FACP SCI_INTERRUPT field, which
///               is going to be used by the ACPI drivers to register
///               sci handler.
/// * `acpi_dev_resource` - resouces needed by the ACPI devices for creating tables
pub fn create_acpi_tables(
    guest_mem: &GuestMemory,
    num_cpus: u8,
    sci_irq: u32,
    acpi_dev_resource: ACPIDevResource,
) -> Option<GuestAddress> {
    // RSDP is at the HI RSDP WINDOW
    let rsdp_offset = GuestAddress(super::ACPI_HI_RSDP_WINDOW_BASE);
    let facs_offset = next_offset(rsdp_offset, RSDP::len() as u64)?;
    let mut offset = next_offset(facs_offset, FACS::len() as u64)?;
    let mut dsdt_offset: Option<GuestAddress> = None;
    let mut tables: Vec<u64> = Vec::new();
    let mut facp: Option<SDT> = None;
    let mut has_madt = false;

    // User supplied System Description Tables, e.g. SSDT.
    for sdt in acpi_dev_resource.sdts.iter() {
        if sdt.is_signature(b"FACP") {
            facp = Some(sdt.clone());
            continue;
        }
        guest_mem.write_at_addr(sdt.as_slice(), offset).ok()?;
        if sdt.is_signature(b"DSDT") {
            dsdt_offset = Some(offset);
        } else {
            tables.push(offset.0);
        }
        if sdt.is_signature(b"APIC") {
            has_madt = true;
        }
        offset = next_offset(offset, sdt.len() as u64)?;
    }

    // FACS
    let facs = FACS::new();
    guest_mem.write_at_addr(facs.as_slice(), facs_offset).ok()?;

    // DSDT
    let dsdt_offset = match dsdt_offset {
        Some(dsdt_offset) => dsdt_offset,
        None => {
            let dsdt_offset = offset;
            let dsdt = create_dsdt_table(acpi_dev_resource.amls);
            guest_mem.write_at_addr(dsdt.as_slice(), offset).ok()?;
            offset = next_offset(offset, dsdt.len() as u64)?;
            dsdt_offset
        }
    };

    // FACP aka FADT
    let pm_iobase = acpi_dev_resource.pm_iobase as u32;
    let mut facp = facp.unwrap_or_else(|| create_facp_table(sci_irq as u16, pm_iobase));

    // Crosvm FACP overrides.
    facp.write(FADT_FIELD_FACS_ADDR32, 0u32);
    facp.write(FADT_FIELD_DSDT_ADDR32, 0u32);
    facp.write(FADT_FIELD_FACS_ADDR, facs_offset.0 as u64);
    facp.write(FADT_FIELD_DSDT_ADDR, dsdt_offset.0 as u64);

    guest_mem.write_at_addr(facp.as_slice(), offset).ok()?;
    tables.push(offset.0);
    offset = next_offset(offset, facp.len() as u64)?;

    // MADT if not provided.
    if !has_madt {
        // MADT
        let mut madt = SDT::new(
            *b"APIC",
            MADT_LEN,
            MADT_REVISION,
            *b"CROSVM",
            *b"CROSVMDT",
            OEM_REVISION,
        );
        madt.write(
            MADT_FIELD_LAPIC_ADDR,
            super::mptable::APIC_DEFAULT_PHYS_BASE as u32,
        );

        for cpu in 0..num_cpus {
            let lapic = LocalAPIC {
                _type: MADT_TYPE_LOCAL_APIC,
                _length: std::mem::size_of::<LocalAPIC>() as u8,
                _processor_id: cpu,
                _apic_id: cpu,
                _flags: MADT_ENABLED,
            };
            madt.append(lapic);
        }

        madt.append(IOAPIC {
            _type: MADT_TYPE_IO_APIC,
            _length: std::mem::size_of::<IOAPIC>() as u8,
            _apic_address: super::mptable::IO_APIC_DEFAULT_PHYS_BASE,
            ..Default::default()
        });

        guest_mem.write_at_addr(madt.as_slice(), offset).ok()?;
        tables.push(offset.0);
        offset = next_offset(offset, madt.len() as u64)?;
    }

    // XSDT
    let mut xsdt = SDT::new(
        *b"XSDT",
        acpi_tables::HEADER_LEN,
        XSDT_REVISION,
        *b"CROSVM",
        *b"CROSVMDT",
        OEM_REVISION,
    );
    for table in tables {
        xsdt.append(table);
    }

    guest_mem.write_at_addr(xsdt.as_slice(), offset).ok()?;

    // RSDP
    let rsdp = RSDP::new(*b"CROSVM", offset.0);
    guest_mem.write_at_addr(rsdp.as_slice(), rsdp_offset).ok()?;

    Some(rsdp_offset)
}
<commit_msg>x86/acpi: clear SMI COMMAND field.<commit_after>// Copyright 2020 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
use acpi_tables::{facs::FACS, rsdp::RSDP, sdt::SDT};
use data_model::DataInit;
use vm_memory::{GuestAddress, GuestMemory};

pub struct ACPIDevResource {
    pub amls: Vec<u8>,
    pub pm_iobase: u64,
    /// Additional system descriptor tables.
    pub sdts: Vec<SDT>,
}

#[repr(C)]
#[derive(Clone, Copy, Default)]
struct LocalAPIC {
    _type: u8,
    _length: u8,
    _processor_id: u8,
    _apic_id: u8,
    _flags: u32,
}

// Safe as LocalAPIC structure only contains raw data
unsafe impl DataInit for LocalAPIC {}

#[repr(C)]
#[derive(Clone, Copy, Default)]
struct IOAPIC {
    _type: u8,
    _length: u8,
    _ioapic_id: u8,
    _reserved: u8,
    _apic_address: u32,
    _gsi_base: u32,
}

// Safe as IOAPIC structure only contains raw data
unsafe impl DataInit for IOAPIC {}

const OEM_REVISION: u32 = 1;
//DSDT
const DSDT_REVISION: u8 = 6;
// FADT
const FADT_LEN: u32 = 276;
const FADT_REVISION: u8 = 6;
const FADT_MINOR_REVISION: u8 = 3;
// FADT flags
const FADT_POWER_BUTTON: u32 = 1 << 4;
const FADT_SLEEP_BUTTON: u32 = 1 << 5;
// FADT fields offset
const FADT_FIELD_FACS_ADDR32: usize = 36;
const FADT_FIELD_DSDT_ADDR32: usize = 40;
const FADT_FIELD_SCI_INTERRUPT: usize = 46;
const FADT_FIELD_SMI_COMMAND: usize = 48;
const FADT_FIELD_PM1A_EVENT_BLK_ADDR: usize = 56;
const FADT_FIELD_PM1A_CONTROL_BLK_ADDR: usize = 64;
const FADT_FIELD_PM1A_EVENT_BLK_LEN: usize = 88;
const FADT_FIELD_PM1A_CONTROL_BLK_LEN: usize = 89;
const FADT_FIELD_FLAGS: usize = 112;
const FADT_FIELD_MINOR_REVISION: usize = 131;
const FADT_FIELD_FACS_ADDR: usize = 132;
const FADT_FIELD_DSDT_ADDR: usize = 140;
const FADT_FIELD_HYPERVISOR_ID: usize = 268;
// MADT
const MADT_LEN: u32 = 44;
const MADT_REVISION: u8 = 5;
// MADT fields offset
const MADT_FIELD_LAPIC_ADDR: usize = 36;
// MADT types
const MADT_TYPE_LOCAL_APIC: u8 = 0;
const MADT_TYPE_IO_APIC: u8 = 1;
// MADT flags
const MADT_ENABLED: u32 = 1;
// XSDT
const XSDT_REVISION: u8 = 1;

fn create_dsdt_table(amls: Vec<u8>) -> SDT {
    let mut dsdt = SDT::new(
        *b"DSDT",
        acpi_tables::HEADER_LEN,
        DSDT_REVISION,
        *b"CROSVM",
        *b"CROSVMDT",
        OEM_REVISION,
    );

    if !amls.is_empty() {
        dsdt.append_slice(amls.as_slice());
    }

    dsdt
}

fn create_facp_table(sci_irq: u16, pm_iobase: u32) -> SDT {
    let mut facp = SDT::new(
        *b"FACP",
        FADT_LEN,
        FADT_REVISION,
        *b"CROSVM",
        *b"CROSVMDT",
        OEM_REVISION,
    );

    let fadt_flags: u32 = FADT_POWER_BUTTON | FADT_SLEEP_BUTTON; // mask POWER and SLEEP BUTTON
    facp.write(FADT_FIELD_FLAGS, fadt_flags);

    // SCI Interrupt
    facp.write(FADT_FIELD_SCI_INTERRUPT, sci_irq);

    // PM1A Event Block Address
    facp.write(FADT_FIELD_PM1A_EVENT_BLK_ADDR, pm_iobase);

    // PM1A Control Block Address
    facp.write(
        FADT_FIELD_PM1A_CONTROL_BLK_ADDR,
        pm_iobase + devices::acpi::ACPIPM_RESOURCE_EVENTBLK_LEN as u32,
    );

    // PM1 Event Block Length
    facp.write(
        FADT_FIELD_PM1A_EVENT_BLK_LEN,
        devices::acpi::ACPIPM_RESOURCE_EVENTBLK_LEN as u8,
    );

    // PM1 Control Block Length
    facp.write(
        FADT_FIELD_PM1A_CONTROL_BLK_LEN,
        devices::acpi::ACPIPM_RESOURCE_CONTROLBLK_LEN as u8,
    );

    facp.write(FADT_FIELD_MINOR_REVISION, FADT_MINOR_REVISION); // FADT minor version
    facp.write(FADT_FIELD_HYPERVISOR_ID, *b"CROSVM"); // Hypervisor Vendor Identity

    facp
}

fn next_offset(offset: GuestAddress, len: u64) -> Option<GuestAddress> {
    // Enforce 64-byte allocation alignment.
    match len % 64 {
        0 => offset.checked_add(len),
        x => offset.checked_add(len.checked_add(64 - x)?),
    }
}

/// Create ACPI tables and return the RSDP.
/// The basic tables DSDT/FACP/MADT/XSDT are constructed in this function.
/// # Arguments
///
/// * `guest_mem` - The guest memory where the tables will be stored.
/// * `num_cpus` - Used to construct the MADT.
/// * `sci_irq` - Used to fill the FACP SCI_INTERRUPT field, which
///               is going to be used by the ACPI drivers to register
///               sci handler.
/// * `acpi_dev_resource` - resouces needed by the ACPI devices for creating tables
pub fn create_acpi_tables(
    guest_mem: &GuestMemory,
    num_cpus: u8,
    sci_irq: u32,
    acpi_dev_resource: ACPIDevResource,
) -> Option<GuestAddress> {
    // RSDP is at the HI RSDP WINDOW
    let rsdp_offset = GuestAddress(super::ACPI_HI_RSDP_WINDOW_BASE);
    let facs_offset = next_offset(rsdp_offset, RSDP::len() as u64)?;
    let mut offset = next_offset(facs_offset, FACS::len() as u64)?;
    let mut dsdt_offset: Option<GuestAddress> = None;
    let mut tables: Vec<u64> = Vec::new();
    let mut facp: Option<SDT> = None;
    let mut has_madt = false;

    // User supplied System Description Tables, e.g. SSDT.
    for sdt in acpi_dev_resource.sdts.iter() {
        if sdt.is_signature(b"FACP") {
            facp = Some(sdt.clone());
            continue;
        }
        guest_mem.write_at_addr(sdt.as_slice(), offset).ok()?;
        if sdt.is_signature(b"DSDT") {
            dsdt_offset = Some(offset);
        } else {
            tables.push(offset.0);
        }
        if sdt.is_signature(b"APIC") {
            has_madt = true;
        }
        offset = next_offset(offset, sdt.len() as u64)?;
    }

    // FACS
    let facs = FACS::new();
    guest_mem.write_at_addr(facs.as_slice(), facs_offset).ok()?;

    // DSDT
    let dsdt_offset = match dsdt_offset {
        Some(dsdt_offset) => dsdt_offset,
        None => {
            let dsdt_offset = offset;
            let dsdt = create_dsdt_table(acpi_dev_resource.amls);
            guest_mem.write_at_addr(dsdt.as_slice(), offset).ok()?;
            offset = next_offset(offset, dsdt.len() as u64)?;
            dsdt_offset
        }
    };

    // FACP aka FADT
    let pm_iobase = acpi_dev_resource.pm_iobase as u32;
    let mut facp = facp.unwrap_or_else(|| create_facp_table(sci_irq as u16, pm_iobase));

    // Crosvm FACP overrides.
    facp.write(FADT_FIELD_SMI_COMMAND, 0u32);
    facp.write(FADT_FIELD_FACS_ADDR32, 0u32);
    facp.write(FADT_FIELD_DSDT_ADDR32, 0u32);
    facp.write(FADT_FIELD_FACS_ADDR, facs_offset.0 as u64);
    facp.write(FADT_FIELD_DSDT_ADDR, dsdt_offset.0 as u64);

    guest_mem.write_at_addr(facp.as_slice(), offset).ok()?;
    tables.push(offset.0);
    offset = next_offset(offset, facp.len() as u64)?;

    // MADT if not provided.
    if !has_madt {
        // MADT
        let mut madt = SDT::new(
            *b"APIC",
            MADT_LEN,
            MADT_REVISION,
            *b"CROSVM",
            *b"CROSVMDT",
            OEM_REVISION,
        );
        madt.write(
            MADT_FIELD_LAPIC_ADDR,
            super::mptable::APIC_DEFAULT_PHYS_BASE as u32,
        );

        for cpu in 0..num_cpus {
            let lapic = LocalAPIC {
                _type: MADT_TYPE_LOCAL_APIC,
                _length: std::mem::size_of::<LocalAPIC>() as u8,
                _processor_id: cpu,
                _apic_id: cpu,
                _flags: MADT_ENABLED,
            };
            madt.append(lapic);
        }

        madt.append(IOAPIC {
            _type: MADT_TYPE_IO_APIC,
            _length: std::mem::size_of::<IOAPIC>() as u8,
            _apic_address: super::mptable::IO_APIC_DEFAULT_PHYS_BASE,
            ..Default::default()
        });

        guest_mem.write_at_addr(madt.as_slice(), offset).ok()?;
        tables.push(offset.0);
        offset = next_offset(offset, madt.len() as u64)?;
    }

    // XSDT
    let mut xsdt = SDT::new(
        *b"XSDT",
        acpi_tables::HEADER_LEN,
        XSDT_REVISION,
        *b"CROSVM",
        *b"CROSVMDT",
        OEM_REVISION,
    );
    for table in tables {
        xsdt.append(table);
    }

    guest_mem.write_at_addr(xsdt.as_slice(), offset).ok()?;

    // RSDP
    let rsdp = RSDP::new(*b"CROSVM", offset.0);
    guest_mem.write_at_addr(rsdp.as_slice(), rsdp_offset).ok()?;

    Some(rsdp_offset)
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::collections::hashmap;
use std::default::Default;
use std::slice;
use core::wire_format;

#[deriving(Show)]
pub enum UnknownValue {
    UnknownFixed32(u32),
    UnknownFixed64(u64),
    UnknownVarint(u64),
    UnknownLengthDelimited(Vec<u8>),
}

impl UnknownValue {
    pub fn wire_type(&self) -> wire_format::WireType {
        self.get_ref().wire_type()
    }

    pub fn get_ref<'s>(&'s self) -> UnknownValueRef<'s> {
        match *self {
            UnknownFixed32(fixed32) => UnknownFixed32Ref(fixed32),
            UnknownFixed64(fixed64) => UnknownFixed64Ref(fixed64),
            UnknownVarint(varint) => UnknownVarintRef(varint),
            UnknownLengthDelimited(ref bytes) => UnknownLengthDelimitedRef(bytes.as_slice()),
        }
    }
}

pub enum UnknownValueRef<'o> {
    UnknownFixed32Ref(u32),
    UnknownFixed64Ref(u64),
    UnknownVarintRef(u64),
    UnknownLengthDelimitedRef(&'o [u8]),
}

impl<'o> UnknownValueRef<'o> {
    pub fn wire_type(&self) -> wire_format::WireType {
        match *self {
            UnknownFixed32Ref(_) => wire_format::WireTypeFixed32,
            UnknownFixed64Ref(_) => wire_format::WireTypeFixed64,
            UnknownVarintRef(_) => wire_format::WireTypeVarint,
            UnknownLengthDelimitedRef(_) => wire_format::WireTypeLengthDelimited,
        }
    }
}

#[deriving(Clone,PartialEq,Eq,Show,Default)]
pub struct UnknownValues {
    pub fixed32: Vec<u32>,
    pub fixed64: Vec<u64>,
    pub varint: Vec<u64>,
    pub length_delimited: Vec<Vec<u8>>,
}

impl UnknownValues {
    pub fn add_value(&mut self, value: UnknownValue) {
        match value {
            UnknownFixed64(fixed64) => self.fixed64.push(fixed64),
            UnknownFixed32(fixed32) => self.fixed32.push(fixed32),
            UnknownVarint(varint) => self.varint.push(varint),
            UnknownLengthDelimited(length_delimited) => self.length_delimited.push(length_delimited),
        };
    }

    pub fn iter<'s>(&'s self) -> UnknownValuesIter<'s> {
        UnknownValuesIter {
            fixed32: self.fixed32.iter(),
            fixed64: self.fixed64.iter(),
            varint: self.varint.iter(),
            length_delimited: self.length_delimited.iter(),
        }
    }
}

pub struct UnknownValuesIter<'o> {
    fixed32: slice::Items<'o, u32>,
    fixed64: slice::Items<'o, u64>,
    varint: slice::Items<'o, u64>,
    length_delimited: slice::Items<'o, Vec<u8>>,
}

impl<'o> Iterator<UnknownValueRef<'o>> for UnknownValuesIter<'o> {
    fn next(&mut self) -> Option<UnknownValueRef<'o>> {
        let fixed32 = self.fixed32.next();
        if fixed32.is_some() {
            return Some(UnknownFixed32Ref(*fixed32.unwrap()));
        }
        let fixed64 = self.fixed64.next();
        if fixed64.is_some() {
            return Some(UnknownFixed64Ref(*fixed64.unwrap()));
        }
        let varint = self.varint.next();
        if varint.is_some() {
            return Some(UnknownVarintRef(*varint.unwrap()));
        }
        let length_delimited = self.length_delimited.next();
        if length_delimited.is_some() {
            return Some(UnknownLengthDelimitedRef(length_delimited.get_ref().as_slice()))
        }
        None
    }
}

#[deriving(Clone,PartialEq,Eq,Show,Default)]
pub struct UnknownFields {
    // option is needed, so it could be placed in static field
    pub fields: Option<HashMap<u32, UnknownValues>>,
}

impl UnknownFields {
    pub fn new() -> UnknownFields {
        Default::default()
    }

    pub fn default_instance() -> &'static UnknownFields {
        static instance: UnknownFields = UnknownFields {
            fields: None,
        };
        &'static instance
    }

    fn init_map(&mut self) {
        if self.fields.is_none() {
            self.fields = Some(Default::default());
        }
    }

    fn find_field<'a>(&'a mut self, number: u32) -> &'a mut UnknownValues {
        self.init_map();

        self.fields.get_mut_ref()
            .find_or_insert_with(number, |_| Default::default())
    }

    pub fn add_fixed32(&mut self, number: u32, fixed32: u32) {
        self.find_field(number).fixed32.push(fixed32);
    }

    pub fn add_fixed64(&mut self, number: u32, fixed64: u64) {
        self.find_field(number).fixed64.push(fixed64);
    }

    pub fn add_varint(&mut self, number: u32, varint: u64) {
        self.find_field(number).varint.push(varint);
    }

    pub fn add_length_delimited(&mut self, number: u32, length_delimited: Vec<u8>) {
        self.find_field(number).length_delimited.push(length_delimited);
    }

    pub fn add_value(&mut self, number: u32, value: UnknownValue) {
        self.find_field(number).add_value(value);
    }

    pub fn iter<'s>(&'s self) -> UnknownFieldIter<'s> {
        UnknownFieldIter {
            entries: self.fields.as_ref().map(|m| m.iter())
        }
    }
}

pub struct UnknownFieldIter<'s> {
    entries: Option<hashmap::Entries<'s, u32, UnknownValues>>,
}

impl<'s> Iterator<(u32, &'s UnknownValues)> for UnknownFieldIter<'s> {
    fn next(&mut self) -> Option<(u32, &'s UnknownValues)> {
        if self.entries.is_none() {
            None
        } else {
            self.entries.get_mut_ref().next().map(|(&number, values)| (number, values))
        }
    }
}
<commit_msg>Drop UnknownFields::default_instance<commit_after>use std::collections::HashMap;
use std::collections::hashmap;
use std::default::Default;
use std::slice;
use core::wire_format;

#[deriving(Show)]
pub enum UnknownValue {
    UnknownFixed32(u32),
    UnknownFixed64(u64),
    UnknownVarint(u64),
    UnknownLengthDelimited(Vec<u8>),
}

impl UnknownValue {
    pub fn wire_type(&self) -> wire_format::WireType {
        self.get_ref().wire_type()
    }

    pub fn get_ref<'s>(&'s self) -> UnknownValueRef<'s> {
        match *self {
            UnknownFixed32(fixed32) => UnknownFixed32Ref(fixed32),
            UnknownFixed64(fixed64) => UnknownFixed64Ref(fixed64),
            UnknownVarint(varint) => UnknownVarintRef(varint),
            UnknownLengthDelimited(ref bytes) => UnknownLengthDelimitedRef(bytes.as_slice()),
        }
    }
}

pub enum UnknownValueRef<'o> {
    UnknownFixed32Ref(u32),
    UnknownFixed64Ref(u64),
    UnknownVarintRef(u64),
    UnknownLengthDelimitedRef(&'o [u8]),
}

impl<'o> UnknownValueRef<'o> {
    pub fn wire_type(&self) -> wire_format::WireType {
        match *self {
            UnknownFixed32Ref(_) => wire_format::WireTypeFixed32,
            UnknownFixed64Ref(_) => wire_format::WireTypeFixed64,
            UnknownVarintRef(_) => wire_format::WireTypeVarint,
            UnknownLengthDelimitedRef(_) => wire_format::WireTypeLengthDelimited,
        }
    }
}

#[deriving(Clone,PartialEq,Eq,Show,Default)]
pub struct UnknownValues {
    pub fixed32: Vec<u32>,
    pub fixed64: Vec<u64>,
    pub varint: Vec<u64>,
    pub length_delimited: Vec<Vec<u8>>,
}

impl UnknownValues {
    pub fn add_value(&mut self, value: UnknownValue) {
        match value {
            UnknownFixed64(fixed64) => self.fixed64.push(fixed64),
            UnknownFixed32(fixed32) => self.fixed32.push(fixed32),
            UnknownVarint(varint) => self.varint.push(varint),
            UnknownLengthDelimited(length_delimited) => self.length_delimited.push(length_delimited),
        };
    }

    pub fn iter<'s>(&'s self) -> UnknownValuesIter<'s> {
        UnknownValuesIter {
            fixed32: self.fixed32.iter(),
            fixed64: self.fixed64.iter(),
            varint: self.varint.iter(),
            length_delimited: self.length_delimited.iter(),
        }
    }
}

pub struct UnknownValuesIter<'o> {
    fixed32: slice::Items<'o, u32>,
    fixed64: slice::Items<'o, u64>,
    varint: slice::Items<'o, u64>,
    length_delimited: slice::Items<'o, Vec<u8>>,
}

impl<'o> Iterator<UnknownValueRef<'o>> for UnknownValuesIter<'o> {
    fn next(&mut self) -> Option<UnknownValueRef<'o>> {
        let fixed32 = self.fixed32.next();
        if fixed32.is_some() {
            return Some(UnknownFixed32Ref(*fixed32.unwrap()));
        }
        let fixed64 = self.fixed64.next();
        if fixed64.is_some() {
            return Some(UnknownFixed64Ref(*fixed64.unwrap()));
        }
        let varint = self.varint.next();
        if varint.is_some() {
            return Some(UnknownVarintRef(*varint.unwrap()));
        }
        let length_delimited = self.length_delimited.next();
        if length_delimited.is_some() {
            return Some(UnknownLengthDelimitedRef(length_delimited.get_ref().as_slice()))
        }
        None
    }
}

#[deriving(Clone,PartialEq,Eq,Show,Default)]
pub struct UnknownFields {
    // option is needed, because HashMap constructor performs allocation,
    // and very expensive
    pub fields: Option<HashMap<u32, UnknownValues>>,
}

impl UnknownFields {
    pub fn new() -> UnknownFields {
        Default::default()
    }

    fn init_map(&mut self) {
        if self.fields.is_none() {
            self.fields = Some(Default::default());
        }
    }

    fn find_field<'a>(&'a mut self, number: u32) -> &'a mut UnknownValues {
        self.init_map();

        self.fields.get_mut_ref()
            .find_or_insert_with(number, |_| Default::default())
    }

    pub fn add_fixed32(&mut self, number: u32, fixed32: u32) {
        self.find_field(number).fixed32.push(fixed32);
    }

    pub fn add_fixed64(&mut self, number: u32, fixed64: u64) {
        self.find_field(number).fixed64.push(fixed64);
    }

    pub fn add_varint(&mut self, number: u32, varint: u64) {
        self.find_field(number).varint.push(varint);
    }

    pub fn add_length_delimited(&mut self, number: u32, length_delimited: Vec<u8>) {
        self.find_field(number).length_delimited.push(length_delimited);
    }

    pub fn add_value(&mut self, number: u32, value: UnknownValue) {
        self.find_field(number).add_value(value);
    }

    pub fn iter<'s>(&'s self) -> UnknownFieldIter<'s> {
        UnknownFieldIter {
            entries: self.fields.as_ref().map(|m| m.iter())
        }
    }
}

pub struct UnknownFieldIter<'s> {
    entries: Option<hashmap::Entries<'s, u32, UnknownValues>>,
}

impl<'s> Iterator<(u32, &'s UnknownValues)> for UnknownFieldIter<'s> {
    fn next(&mut self) -> Option<(u32, &'s UnknownValues)> {
        if self.entries.is_none() {
            None
        } else {
            self.entries.get_mut_ref().next().map(|(&number, values)| (number, values))
        }
    }
}
<|endoftext|>


<commit_before><commit_msg>introduce NO_PRIORITIES, which describes the maximum number of priorities<commit_after><|endoftext|>


<commit_before>//! Support code for rustc's built in unit-test and micro-benchmarking
//! framework.
//!
//! Almost all user code will only be interested in `Bencher` and
//! `black_box`. All other interactions (such as writing tests and
//! benchmarks themselves) should be done via the `#[test]` and
//! `#[bench]` attributes.
//!
//! See the [Testing Chapter](../book/ch11-00-testing.html) of the book for more details.

#![crate_name = "test"]
#![unstable(feature = "test", issue = "27812")]
#![doc(html_root_url = "https://doc.rust-lang.org/nightly/",
       test(attr(deny(warnings))))]
#![feature(test)]

extern crate libtest;
pub use libtest::*;
<commit_msg>Allow the staged_api<commit_after>//! Support code for rustc's built in unit-test and micro-benchmarking
//! framework.
//!
//! Almost all user code will only be interested in `Bencher` and
//! `black_box`. All other interactions (such as writing tests and
//! benchmarks themselves) should be done via the `#[test]` and
//! `#[bench]` attributes.
//!
//! See the [Testing Chapter](../book/ch11-00-testing.html) of the book for more details.

#![crate_name = "test"]
#![unstable(feature = "test", issue = "27812")]
#![doc(html_root_url = "https://doc.rust-lang.org/nightly/",
       test(attr(deny(warnings))))]
#![feature(staged_api)]
#![feature(test)]

extern crate libtest;
pub use libtest::*;
<|endoftext|>


<commit_before>// xfail-fast  (compile-flags unsupported on windows)
// compile-flags:--borrowck=err
// exec-env:RUST_POISON_ON_FREE=1

fn main() {
    let mut x = @mut @{f: ~3};
    alt x {
      @@{f: b_x} {
        assert *b_x == 3;
        assert ptr::addr_of(x.f) == ptr::addr_of(b_x);

        *x = @{f: ~4};

        #debug["ptr::addr_of(*b_x) = %x", ptr::addr_of(*b_x) as uint];
        assert *b_x == 3;
        assert ptr::addr_of(*x.f) != ptr::addr_of(*b_x);
      }
    }
}<commit_msg>xfail-test failing test<commit_after>// xfail-test  (still buggy)
// xfail-fast  (compile-flags unsupported on windows)
// compile-flags:--borrowck=err
// exec-env:RUST_POISON_ON_FREE=1

fn main() {
    let mut x = @mut @{f: ~3};
    alt x {
      @@{f: b_x} {
        assert *b_x == 3;
        assert ptr::addr_of(x.f) == ptr::addr_of(b_x);

        *x = @{f: ~4};

        #debug["ptr::addr_of(*b_x) = %x", ptr::addr_of(*b_x) as uint];
        assert *b_x == 3;
        assert ptr::addr_of(*x.f) != ptr::addr_of(*b_x);
      }
    }
}
<|endoftext|>


<commit_before>#[cfg(test)]
mod tests;
pub mod response;

pub use message::response::*;
use notification::Notification;
use std::collections::HashMap;
use std::collections::BTreeMap;
use std::str;
use std::io::prelude::*;
use hyper::Client;
use hyper::header::{Authorization, ContentType};
use hyper::status::StatusCode;
use rustc_serialize::json::{self, Json, ToJson};

#[derive(PartialEq, Debug, Clone)]
pub enum Priority {
    Normal, High
}

/// Represents a FCM message. Construct the FCM message
/// using various utility methods and finally send it.
/// # Examples:
/// ```rust
/// use fcm::Message;
///
/// let message = Message::new("<registration id>").dry_run(true);
/// ```
pub struct Message<'a> {
    to: &'a str,
    registration_ids: Option<Vec<String>>,
    collapse_key: Option<&'a str>,
    priority: Option<Priority>,
    content_available: Option<bool>,
    delay_while_idle: Option<bool>,
    time_to_live: Option<i32>,
    restricted_package_name: Option<&'a str>,
    dry_run: Option<bool>,
    data: Option<HashMap<String, String>>,
    notification: Option<Notification<'a>>,
}

impl <'a> ToJson for Message<'a> {
    fn to_json(&self) -> Json {
        let mut root = BTreeMap::new();

        root.insert("to".to_string(), self.to.to_json());

        if self.registration_ids.is_some() {
            root.insert("registration_ids".to_string(),
                    self.registration_ids.clone().unwrap().to_json());
        }

        if self.collapse_key.is_some() {
            root.insert("collapse_key".to_string(), self.collapse_key.clone().unwrap().to_json());
        }

        if self.priority.is_some() {
            root.insert("priority".to_string(), match self.priority.clone().unwrap() {
                Priority::Normal => Json::String("normal".to_string()),
                Priority::High => Json::String("high".to_string()),
            });
        }

        if self.content_available.is_some() {
            root.insert("content_available".to_string(), self.content_available.unwrap().to_json());
        }

        if self.delay_while_idle.is_some() {
            root.insert("delay_while_idle".to_string(), self.delay_while_idle.unwrap().to_json());
        }

        if self.time_to_live.is_some() {
            root.insert("time_to_live".to_string(), self.time_to_live.unwrap().to_json());
        }

        if self.restricted_package_name.is_some() {
            root.insert("restricted_package_name".to_string(), self.restricted_package_name.clone().unwrap().to_json());
        }

        if self.dry_run.is_some() {
            root.insert("dry_run".to_string(), self.dry_run.unwrap().to_json());
        }

        if self.data.is_some() {
            root.insert("data".to_string(), self.data.clone().unwrap().to_json());
        }

        if self.notification.is_some() {
            root.insert("notification".to_string(), self.notification.clone().unwrap().to_json());
        }

        Json::Object(root)
    }
}

impl <'a> Message<'a> {
    /// Get a new instance of Message. You need to supply either
    /// a registration id, or a topic (/topic/...).
    pub fn new(to: &'a str) -> Message {
        Message {
            to: to,
            registration_ids: None,
            collapse_key: None,
            priority: None,
            content_available: None,
            delay_while_idle: None,
            time_to_live: None,
            restricted_package_name: None,
            dry_run: None,
            data: None,
            notification: None,
        }
    }

    /// Set various registration ids to which the message ought to be sent.
    pub fn registration_ids(mut self, ids: Vec<&'a str>) -> Message<'a> {
        self.registration_ids = Some(ids.iter().map(|s| s.to_string()).collect());
        self
    }

    /// Set this parameter to identify groups of messages that can be collapsed.
    pub fn collapse_key(mut self, collapse_key: &'a str) -> Message<'a> {
        self.collapse_key = Some(collapse_key);
        self
    }

    /// Set the priority of the message. You can set Normal or High priorities.
    /// # Examples:
    /// ```rust
    /// use fcm::{Message, Priority};
    ///
    /// let message = Message::new("<registration id>")
    ///     .priority(Priority::High);
    /// ```
    pub fn priority(mut self, priority: Priority) -> Message<'a> {
        self.priority = Some(priority);
        self
    }

    /// To set the `content-available` field on iOS
    pub fn content_available(mut self, content_available: bool) -> Message<'a> {
        self.content_available = Some(content_available);
        self
    }

    /// When set to `true`, sends the message only when the device is active.
    pub fn delay_while_idle(mut self, delay_while_idle: bool) -> Message<'a> {
        self.delay_while_idle = Some(delay_while_idle);
        self
    }

    /// How long (in seconds) to keep the message on FCM servers in case the device
    /// is offline. The maximum and default is 4 weeks.
    pub fn time_to_live(mut self, time_to_live: i32) -> Message<'a> {
        self.time_to_live = Some(time_to_live);
        self
    }

    /// Package name of the application where the registration tokens must match.
    pub fn restricted_package_name(mut self, restricted_package_name: &'a str) -> Message<'a> {
        self.restricted_package_name = Some(restricted_package_name);
        self
    }

    /// When set to `true`, allows you to test FCM without actually sending the message.
    pub fn dry_run(mut self, dry_run: bool) -> Message<'a> {
        self.dry_run = Some(dry_run);
        self
    }

    /// Use this to add custom key-value pairs to the message. This data
    /// must be handled appropriately on the client end.
    /// # Examples:
    /// ```rust
    /// use fcm::Message;
    /// use std::collections::HashMap;
    ///
    /// let mut map = HashMap::new();
    /// map.insert("message", "Howdy!");
    ///
    /// let message = Message::new("<registration id>").data(map);
    /// ```
    pub fn data(mut self, data: HashMap<&'a str, &'a str>) -> Message<'a> {
        let mut datamap: HashMap<String, String> = HashMap::new();
        for (key, val) in data.iter() {
            datamap.insert(key.to_string(), val.to_string());
        }

        self.data = Some(datamap);
        self
    }

    /// Use this to set a `Notification` for the message.
    /// # Examples:
    /// ```rust
    /// use fcm::{Message, NotificationBuilder};
    ///
    /// let notification = NotificationBuilder::new("Hey!")
    ///     .body("Do you want to catch up later?")
    ///     .finalize();
    ///
    /// let message = Message::new("<registration id>")
    ///     .notification(notification);
    /// ```
    pub fn notification(mut self, notification: Notification<'a>) -> Message<'a> {
        self.notification = Some(notification);
        self
    }

    /// Send the message using your FCM API Key.
    /// # Examples:
    /// ```no_run
    /// use fcm::Message;
    /// use std::collections::HashMap;
    ///
    /// let mut map = HashMap::new();
    /// map.insert("message", "Howdy!");
    ///
    /// let result = Message::new("<registration id>")
    ///     .data(map)
    ///     .send("<FCM API Key>");
    /// ```
    pub fn send(self, api_key: &'a str) -> Result<response::FcmResponse, response::FcmError> {
        let payload = self.to_json().to_string();
        let auth_header = "key=".to_string() + api_key;
        let client = Client::new();

        let response = client
            .post("https://fcm.googleapis.com/fcm/send")
            .body(&payload)
            .header(Authorization(auth_header))
            .header(ContentType::json())
            .send();

        match response {
            Ok(mut response) => {
                let mut body = String::new();
                response.read_to_string(&mut body).unwrap();

                Message::parse_response(response.status, &body)
            },
            Err(_) => {
                Message::parse_response(StatusCode::InternalServerError, "Server Error")
            }
        }
    }

    fn parse_response(status: StatusCode, body: &str) -> Result<response::FcmResponse, response::FcmError> {
        use hyper::status::StatusCode::*;
        use hyper::status::StatusClass::ServerError;
        match status {
            Ok => Result::Ok(json::decode(body).unwrap()),
            Unauthorized => Err(response::FcmError::Unauthorized),
            BadRequest => Err(response::FcmError::InvalidMessage(body.to_string())),
            _ => match status.class() {
                ServerError => Err(response::FcmError::ServerError),
                _ => Err(response::FcmError::InvalidMessage("Unknown Error".to_string())),
            }
        }
    }
}
<commit_msg>Remove unused import<commit_after>#[cfg(test)]
mod tests;
pub mod response;

pub use message::response::*;
use notification::Notification;
use std::collections::HashMap;
use std::collections::BTreeMap;
use std::io::prelude::*;
use hyper::Client;
use hyper::header::{Authorization, ContentType};
use hyper::status::StatusCode;
use rustc_serialize::json::{self, Json, ToJson};

#[derive(PartialEq, Debug, Clone)]
pub enum Priority {
    Normal, High
}

/// Represents a FCM message. Construct the FCM message
/// using various utility methods and finally send it.
/// # Examples:
/// ```rust
/// use fcm::Message;
///
/// let message = Message::new("<registration id>").dry_run(true);
/// ```
pub struct Message<'a> {
    to: &'a str,
    registration_ids: Option<Vec<String>>,
    collapse_key: Option<&'a str>,
    priority: Option<Priority>,
    content_available: Option<bool>,
    delay_while_idle: Option<bool>,
    time_to_live: Option<i32>,
    restricted_package_name: Option<&'a str>,
    dry_run: Option<bool>,
    data: Option<HashMap<String, String>>,
    notification: Option<Notification<'a>>,
}

impl <'a> ToJson for Message<'a> {
    fn to_json(&self) -> Json {
        let mut root = BTreeMap::new();

        root.insert("to".to_string(), self.to.to_json());

        if self.registration_ids.is_some() {
            root.insert("registration_ids".to_string(),
                    self.registration_ids.clone().unwrap().to_json());
        }

        if self.collapse_key.is_some() {
            root.insert("collapse_key".to_string(), self.collapse_key.clone().unwrap().to_json());
        }

        if self.priority.is_some() {
            root.insert("priority".to_string(), match self.priority.clone().unwrap() {
                Priority::Normal => Json::String("normal".to_string()),
                Priority::High => Json::String("high".to_string()),
            });
        }

        if self.content_available.is_some() {
            root.insert("content_available".to_string(), self.content_available.unwrap().to_json());
        }

        if self.delay_while_idle.is_some() {
            root.insert("delay_while_idle".to_string(), self.delay_while_idle.unwrap().to_json());
        }

        if self.time_to_live.is_some() {
            root.insert("time_to_live".to_string(), self.time_to_live.unwrap().to_json());
        }

        if self.restricted_package_name.is_some() {
            root.insert("restricted_package_name".to_string(), self.restricted_package_name.clone().unwrap().to_json());
        }

        if self.dry_run.is_some() {
            root.insert("dry_run".to_string(), self.dry_run.unwrap().to_json());
        }

        if self.data.is_some() {
            root.insert("data".to_string(), self.data.clone().unwrap().to_json());
        }

        if self.notification.is_some() {
            root.insert("notification".to_string(), self.notification.clone().unwrap().to_json());
        }

        Json::Object(root)
    }
}

impl <'a> Message<'a> {
    /// Get a new instance of Message. You need to supply either
    /// a registration id, or a topic (/topic/...).
    pub fn new(to: &'a str) -> Message {
        Message {
            to: to,
            registration_ids: None,
            collapse_key: None,
            priority: None,
            content_available: None,
            delay_while_idle: None,
            time_to_live: None,
            restricted_package_name: None,
            dry_run: None,
            data: None,
            notification: None,
        }
    }

    /// Set various registration ids to which the message ought to be sent.
    pub fn registration_ids(mut self, ids: Vec<&'a str>) -> Message<'a> {
        self.registration_ids = Some(ids.iter().map(|s| s.to_string()).collect());
        self
    }

    /// Set this parameter to identify groups of messages that can be collapsed.
    pub fn collapse_key(mut self, collapse_key: &'a str) -> Message<'a> {
        self.collapse_key = Some(collapse_key);
        self
    }

    /// Set the priority of the message. You can set Normal or High priorities.
    /// # Examples:
    /// ```rust
    /// use fcm::{Message, Priority};
    ///
    /// let message = Message::new("<registration id>")
    ///     .priority(Priority::High);
    /// ```
    pub fn priority(mut self, priority: Priority) -> Message<'a> {
        self.priority = Some(priority);
        self
    }

    /// To set the `content-available` field on iOS
    pub fn content_available(mut self, content_available: bool) -> Message<'a> {
        self.content_available = Some(content_available);
        self
    }

    /// When set to `true`, sends the message only when the device is active.
    pub fn delay_while_idle(mut self, delay_while_idle: bool) -> Message<'a> {
        self.delay_while_idle = Some(delay_while_idle);
        self
    }

    /// How long (in seconds) to keep the message on FCM servers in case the device
    /// is offline. The maximum and default is 4 weeks.
    pub fn time_to_live(mut self, time_to_live: i32) -> Message<'a> {
        self.time_to_live = Some(time_to_live);
        self
    }

    /// Package name of the application where the registration tokens must match.
    pub fn restricted_package_name(mut self, restricted_package_name: &'a str) -> Message<'a> {
        self.restricted_package_name = Some(restricted_package_name);
        self
    }

    /// When set to `true`, allows you to test FCM without actually sending the message.
    pub fn dry_run(mut self, dry_run: bool) -> Message<'a> {
        self.dry_run = Some(dry_run);
        self
    }

    /// Use this to add custom key-value pairs to the message. This data
    /// must be handled appropriately on the client end.
    /// # Examples:
    /// ```rust
    /// use fcm::Message;
    /// use std::collections::HashMap;
    ///
    /// let mut map = HashMap::new();
    /// map.insert("message", "Howdy!");
    ///
    /// let message = Message::new("<registration id>").data(map);
    /// ```
    pub fn data(mut self, data: HashMap<&'a str, &'a str>) -> Message<'a> {
        let mut datamap: HashMap<String, String> = HashMap::new();
        for (key, val) in data.iter() {
            datamap.insert(key.to_string(), val.to_string());
        }

        self.data = Some(datamap);
        self
    }

    /// Use this to set a `Notification` for the message.
    /// # Examples:
    /// ```rust
    /// use fcm::{Message, NotificationBuilder};
    ///
    /// let notification = NotificationBuilder::new("Hey!")
    ///     .body("Do you want to catch up later?")
    ///     .finalize();
    ///
    /// let message = Message::new("<registration id>")
    ///     .notification(notification);
    /// ```
    pub fn notification(mut self, notification: Notification<'a>) -> Message<'a> {
        self.notification = Some(notification);
        self
    }

    /// Send the message using your FCM API Key.
    /// # Examples:
    /// ```no_run
    /// use fcm::Message;
    /// use std::collections::HashMap;
    ///
    /// let mut map = HashMap::new();
    /// map.insert("message", "Howdy!");
    ///
    /// let result = Message::new("<registration id>")
    ///     .data(map)
    ///     .send("<FCM API Key>");
    /// ```
    pub fn send(self, api_key: &'a str) -> Result<response::FcmResponse, response::FcmError> {
        let payload = self.to_json().to_string();
        let auth_header = "key=".to_string() + api_key;
        let client = Client::new();

        let response = client
            .post("https://fcm.googleapis.com/fcm/send")
            .body(&payload)
            .header(Authorization(auth_header))
            .header(ContentType::json())
            .send();

        match response {
            Ok(mut response) => {
                let mut body = String::new();
                response.read_to_string(&mut body).unwrap();

                Message::parse_response(response.status, &body)
            },
            Err(_) => {
                Message::parse_response(StatusCode::InternalServerError, "Server Error")
            }
        }
    }

    fn parse_response(status: StatusCode, body: &str) -> Result<response::FcmResponse, response::FcmError> {
        use hyper::status::StatusCode::*;
        use hyper::status::StatusClass::ServerError;
        match status {
            Ok => Result::Ok(json::decode(body).unwrap()),
            Unauthorized => Err(response::FcmError::Unauthorized),
            BadRequest => Err(response::FcmError::InvalidMessage(body.to_string())),
            _ => match status.class() {
                ServerError => Err(response::FcmError::ServerError),
                _ => Err(response::FcmError::InvalidMessage("Unknown Error".to_string())),
            }
        }
    }
}
<|endoftext|>


<commit_before>use bytes::Buf;
use std::collections::vec_deque;
use std::collections::VecDeque;
use std::io::IoSlice;
use std::mem;
use std::ops::Deref;
use std::ops::DerefMut;

#[derive(Debug)]
pub(crate) struct BufVecDeque<B: Buf> {
    deque: VecDeque<B>,
    len: usize,
}

impl<B: Buf> Default for BufVecDeque<B> {
    fn default() -> Self {
        BufVecDeque {
            deque: VecDeque::default(),
            len: 0,
        }
    }
}

impl<B: Buf, I: Into<VecDeque<B>>> From<I> for BufVecDeque<B> {
    fn from(deque: I) -> Self {
        let deque = deque.into();
        let len = deque.iter().map(Buf::remaining).sum();
        BufVecDeque { deque, len }
    }
}

impl<B: Buf> BufVecDeque<B> {
    pub fn new() -> BufVecDeque<B> {
        Default::default()
    }

    pub fn len(&self) -> usize {
        self.len
    }

    pub fn push_back(&mut self, bytes: B) {
        self.len += bytes.remaining();
        self.deque.push_back(bytes);
    }

    pub fn back_mut(&mut self) -> Option<BufVecDequeBackMut<B>> {
        match self.deque.pop_back() {
            Some(back) => Some(BufVecDequeBackMut {
                deque: self,
                remaining: back.remaining(),
                back: Some(back),
            }),
            None => None,
        }
    }
}

impl<B: Buf> Buf for BufVecDeque<B> {
    fn remaining(&self) -> usize {
        self.len
    }

    fn bytes(&self) -> &[u8] {
        for b in &self.deque {
            let bytes = b.bytes();
            if !bytes.is_empty() {
                return bytes;
            }
        }
        &[]
    }

    fn bytes_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize {
        let mut n = 0;
        for b in &self.deque {
            if n == dst.len() {
                break;
            }
            n += b.bytes_vectored(&mut dst[n..]);
        }
        n
    }

    fn advance(&mut self, mut cnt: usize) {
        assert!(self.len >= cnt);
        self.len -= cnt;

        while cnt != 0 {
            let front = self.deque.front_mut().unwrap();
            let front_remaining = front.remaining();
            if cnt < front_remaining {
                front.advance(cnt);
                break;
            }

            self.deque.pop_front().unwrap();

            cnt -= front_remaining;
        }
    }
}

impl<B: Buf> IntoIterator for BufVecDeque<B> {
    type Item = B;
    type IntoIter = vec_deque::IntoIter<B>;

    fn into_iter(self) -> Self::IntoIter {
        self.deque.into_iter()
    }
}

pub struct BufVecDequeBackMut<'a, B: Buf> {
    deque: &'a mut BufVecDeque<B>,
    back: Option<B>,
    remaining: usize,
}

impl<'a, B: Buf> Deref for BufVecDequeBackMut<'a, B> {
    type Target = B;

    fn deref(&self) -> &Self::Target {
        self.back.as_ref().unwrap()
    }
}

impl<'a, B: Buf> DerefMut for BufVecDequeBackMut<'a, B> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.back.as_mut().unwrap()
    }
}

impl<'a, B: Buf> Drop for BufVecDequeBackMut<'a, B> {
    fn drop(&mut self) {
        let back = mem::take(&mut self.back).unwrap();
        let new_remaining = back.remaining();
        if new_remaining > self.remaining {
            self.deque.len += new_remaining - self.remaining;
        } else {
            self.deque.len -= self.remaining - new_remaining;
        }
        self.deque.deque.push_back(back);
    }
}

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

    #[test]
    fn back_mut() {
        let mut d = BufVecDeque::<VecDeque<u8>>::new();
        d.push_back(VecDeque::from(vec![3, 4]));
        d.push_back(VecDeque::from(vec![4, 6]));
        assert_eq!(4, d.remaining());
        d.back_mut().unwrap().push_back(7);
        assert_eq!(5, d.remaining());
        d.back_mut().unwrap().pop_back();
        assert_eq!(4, d.remaining());
        d.back_mut().unwrap().pop_back();
        d.back_mut().unwrap().pop_back();
        assert_eq!(2, d.remaining());
    }
}
<commit_msg>BufVecDeque::pop_back<commit_after>use bytes::Buf;
use std::collections::vec_deque;
use std::collections::VecDeque;
use std::io::IoSlice;
use std::mem;
use std::ops::Deref;
use std::ops::DerefMut;

#[derive(Debug)]
pub(crate) struct BufVecDeque<B: Buf> {
    deque: VecDeque<B>,
    len: usize,
}

impl<B: Buf> Default for BufVecDeque<B> {
    fn default() -> Self {
        BufVecDeque {
            deque: VecDeque::default(),
            len: 0,
        }
    }
}

impl<B: Buf, I: Into<VecDeque<B>>> From<I> for BufVecDeque<B> {
    fn from(deque: I) -> Self {
        let deque = deque.into();
        let len = deque.iter().map(Buf::remaining).sum();
        BufVecDeque { deque, len }
    }
}

impl<B: Buf> BufVecDeque<B> {
    pub fn new() -> BufVecDeque<B> {
        Default::default()
    }

    pub fn len(&self) -> usize {
        self.len
    }

    pub fn push_back(&mut self, bytes: B) {
        self.len += bytes.remaining();
        self.deque.push_back(bytes);
    }

    pub fn pop_back(&mut self) -> Option<B> {
        match self.deque.pop_back() {
            Some(b) => {
                self.len -= b.remaining();
                Some(b)
            }
            None => None,
        }
    }

    pub fn back_mut(&mut self) -> Option<BufVecDequeBackMut<B>> {
        match self.deque.pop_back() {
            Some(back) => Some(BufVecDequeBackMut {
                deque: self,
                remaining: back.remaining(),
                back: Some(back),
            }),
            None => None,
        }
    }
}

impl<B: Buf> Buf for BufVecDeque<B> {
    fn remaining(&self) -> usize {
        self.len
    }

    fn bytes(&self) -> &[u8] {
        for b in &self.deque {
            let bytes = b.bytes();
            if !bytes.is_empty() {
                return bytes;
            }
        }
        &[]
    }

    fn bytes_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize {
        let mut n = 0;
        for b in &self.deque {
            if n == dst.len() {
                break;
            }
            n += b.bytes_vectored(&mut dst[n..]);
        }
        n
    }

    fn advance(&mut self, mut cnt: usize) {
        assert!(self.len >= cnt);
        self.len -= cnt;

        while cnt != 0 {
            let front = self.deque.front_mut().unwrap();
            let front_remaining = front.remaining();
            if cnt < front_remaining {
                front.advance(cnt);
                break;
            }

            self.deque.pop_front().unwrap();

            cnt -= front_remaining;
        }
    }
}

impl<B: Buf> IntoIterator for BufVecDeque<B> {
    type Item = B;
    type IntoIter = vec_deque::IntoIter<B>;

    fn into_iter(self) -> Self::IntoIter {
        self.deque.into_iter()
    }
}

pub struct BufVecDequeBackMut<'a, B: Buf> {
    deque: &'a mut BufVecDeque<B>,
    back: Option<B>,
    remaining: usize,
}

impl<'a, B: Buf> Deref for BufVecDequeBackMut<'a, B> {
    type Target = B;

    fn deref(&self) -> &Self::Target {
        self.back.as_ref().unwrap()
    }
}

impl<'a, B: Buf> DerefMut for BufVecDequeBackMut<'a, B> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.back.as_mut().unwrap()
    }
}

impl<'a, B: Buf> Drop for BufVecDequeBackMut<'a, B> {
    fn drop(&mut self) {
        let back = mem::take(&mut self.back).unwrap();
        let new_remaining = back.remaining();
        if new_remaining > self.remaining {
            self.deque.len += new_remaining - self.remaining;
        } else {
            self.deque.len -= self.remaining - new_remaining;
        }
        self.deque.deque.push_back(back);
    }
}

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

    #[test]
    fn back_mut() {
        let mut d = BufVecDeque::<VecDeque<u8>>::new();
        d.push_back(VecDeque::from(vec![3, 4]));
        d.push_back(VecDeque::from(vec![4, 6]));
        assert_eq!(4, d.remaining());
        d.back_mut().unwrap().push_back(7);
        assert_eq!(5, d.remaining());
        d.back_mut().unwrap().pop_back();
        assert_eq!(4, d.remaining());
        d.back_mut().unwrap().pop_back();
        d.back_mut().unwrap().pop_back();
        assert_eq!(2, d.remaining());
    }

    #[test]
    fn pop_back() {
        let mut d = BufVecDeque::<VecDeque<u8>>::new();
        d.push_back(VecDeque::from(vec![3, 4]));
        d.push_back(VecDeque::from(vec![4, 6, 7]));

        d.pop_back().unwrap();
        assert_eq!(2, d.remaining());
    }
}
<|endoftext|>


<commit_before>use crate::bytes_ext::buf_vec_deque::BufVecDeque;
use bytes::Buf;
use bytes::Bytes;
use bytes::BytesMut;
use std::collections::VecDeque;
use std::io::IoSlice;

#[derive(Debug, Default)]
pub(crate) struct BytesVecDeque {
    deque: BufVecDeque<Bytes>,
}

impl<I: Into<VecDeque<Bytes>>> From<I> for BytesVecDeque {
    fn from(deque: I) -> Self {
        BytesVecDeque {
            deque: deque.into().into(),
        }
    }
}

impl Into<Bytes> for BytesVecDeque {
    fn into(self) -> Bytes {
        self.into_bytes()
    }
}

impl Into<Vec<u8>> for BytesVecDeque {
    fn into(self) -> Vec<u8> {
        let mut v = Vec::with_capacity(self.remaining());
        for b in self.deque {
            v.extend_from_slice(b.as_ref());
        }
        v
    }
}

impl BytesVecDeque {
    pub fn new() -> BytesVecDeque {
        Default::default()
    }

    pub fn len(&self) -> usize {
        self.deque.remaining()
    }

    pub fn extend(&mut self, bytes: Bytes) {
        if bytes.is_empty() {
            return;
        }
        self.deque.push_back(bytes);
    }

    pub fn get_bytes(&self) -> Bytes {
        let mut bytes_mut = BytesMut::with_capacity(self.remaining());
        for b in &self.deque {
            bytes_mut.extend_from_slice(b);
        }
        bytes_mut.freeze()
    }

    pub fn into_bytes(mut self) -> Bytes {
        self.to_bytes()
    }
}

impl Buf for BytesVecDeque {
    fn remaining(&self) -> usize {
        self.deque.remaining()
    }

    fn bytes(&self) -> &[u8] {
        self.deque.bytes()
    }

    fn bytes_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize {
        self.deque.bytes_vectored(dst)
    }

    fn advance(&mut self, cnt: usize) {
        self.deque.advance(cnt)
    }

    fn to_bytes(&mut self) -> Bytes {
        self.deque.to_bytes()
    }
}

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

    #[test]
    fn buf_empty() {
        let d = BytesVecDeque::new();
        assert_eq!(&[0u8; 0], Buf::bytes(&d));
        assert_eq!(0, Buf::remaining(&d));
        assert_eq!(false, Buf::has_remaining(&d));
    }

    #[test]
    fn buf_advance_full() {
        let mut d = BytesVecDeque::new();
        d.extend(Bytes::from_static(b"ab"));
        d.extend(Bytes::from_static(b"cde"));

        assert_eq!(b"ab", Buf::bytes(&d));
        Buf::advance(&mut d, 2);
        assert_eq!(b"cde", Buf::bytes(&d));
        Buf::advance(&mut d, 3);
        assert_eq!(0, Buf::remaining(&d));
        assert_eq!(false, Buf::has_remaining(&d));
    }

    #[test]
    fn buf_advance() {
        let mut d = BytesVecDeque::new();
        d.extend(Bytes::from_static(b"ab"));
        d.extend(Bytes::from_static(b"cde"));

        assert_eq!(b"ab", Buf::bytes(&d));
        Buf::advance(&mut d, 1);
        assert_eq!(b"b", Buf::bytes(&d));
        Buf::advance(&mut d, 3);
        assert_eq!(b"e", Buf::bytes(&d));
        Buf::advance(&mut d, 1);
        assert_eq!(0, Buf::remaining(&d));
        assert_eq!(false, Buf::has_remaining(&d));
    }

    #[test]
    fn buf_bytes_vectored() {
        let mut d = BytesVecDeque::new();
        d.extend(Bytes::from_static(b"ab"));
        d.extend(Bytes::from_static(b"cde"));

        let mut v = [IoSlice::new(&[])];
        assert_eq!(1, d.bytes_vectored(&mut v));
        assert_eq!(b"ab", &*v[0]);

        let mut v = [IoSlice::new(&[]), IoSlice::new(&[])];
        assert_eq!(2, d.bytes_vectored(&mut v));
        assert_eq!(b"ab", &*v[0]);
        assert_eq!(b"cde", &*v[1]);

        let mut v = [IoSlice::new(&[]), IoSlice::new(&[]), IoSlice::new(&[])];
        assert_eq!(2, d.bytes_vectored(&mut v));
        assert_eq!(b"ab", &*v[0]);
        assert_eq!(b"cde", &*v[1]);
    }
}
<commit_msg>BytesVecDeque::get_bytes<commit_after>use crate::bytes_ext::buf_vec_deque::BufVecDeque;
use crate::BufGetBytes;
use bytes::Buf;
use bytes::Bytes;
use bytes::BytesMut;
use std::collections::VecDeque;
use std::io::IoSlice;

#[derive(Debug, Default)]
pub(crate) struct BytesVecDeque {
    deque: BufVecDeque<Bytes>,
}

impl<I: Into<VecDeque<Bytes>>> From<I> for BytesVecDeque {
    fn from(deque: I) -> Self {
        BytesVecDeque {
            deque: deque.into().into(),
        }
    }
}

impl Into<Bytes> for BytesVecDeque {
    fn into(self) -> Bytes {
        self.into_bytes()
    }
}

impl Into<Vec<u8>> for BytesVecDeque {
    fn into(self) -> Vec<u8> {
        let mut v = Vec::with_capacity(self.remaining());
        for b in self.deque {
            v.extend_from_slice(b.as_ref());
        }
        v
    }
}

impl BytesVecDeque {
    pub fn new() -> BytesVecDeque {
        Default::default()
    }

    pub fn len(&self) -> usize {
        self.deque.remaining()
    }

    pub fn extend(&mut self, bytes: Bytes) {
        if bytes.is_empty() {
            return;
        }
        self.deque.push_back(bytes);
    }

    pub fn get_bytes(&self) -> Bytes {
        let mut bytes_mut = BytesMut::with_capacity(self.remaining());
        for b in &self.deque {
            bytes_mut.extend_from_slice(b);
        }
        bytes_mut.freeze()
    }

    pub fn into_bytes(mut self) -> Bytes {
        self.to_bytes()
    }
}

impl Buf for BytesVecDeque {
    fn remaining(&self) -> usize {
        self.deque.remaining()
    }

    fn bytes(&self) -> &[u8] {
        self.deque.bytes()
    }

    fn bytes_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize {
        self.deque.bytes_vectored(dst)
    }

    fn advance(&mut self, cnt: usize) {
        self.deque.advance(cnt)
    }

    fn to_bytes(&mut self) -> Bytes {
        self.deque.to_bytes()
    }
}

impl BufGetBytes for BytesVecDeque {
    fn get_bytes(&mut self, cnt: usize) -> Bytes {
        self.deque.get_bytes(cnt)
    }
}

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

    #[test]
    fn buf_empty() {
        let d = BytesVecDeque::new();
        assert_eq!(&[0u8; 0], Buf::bytes(&d));
        assert_eq!(0, Buf::remaining(&d));
        assert_eq!(false, Buf::has_remaining(&d));
    }

    #[test]
    fn buf_advance_full() {
        let mut d = BytesVecDeque::new();
        d.extend(Bytes::from_static(b"ab"));
        d.extend(Bytes::from_static(b"cde"));

        assert_eq!(b"ab", Buf::bytes(&d));
        Buf::advance(&mut d, 2);
        assert_eq!(b"cde", Buf::bytes(&d));
        Buf::advance(&mut d, 3);
        assert_eq!(0, Buf::remaining(&d));
        assert_eq!(false, Buf::has_remaining(&d));
    }

    #[test]
    fn buf_advance() {
        let mut d = BytesVecDeque::new();
        d.extend(Bytes::from_static(b"ab"));
        d.extend(Bytes::from_static(b"cde"));

        assert_eq!(b"ab", Buf::bytes(&d));
        Buf::advance(&mut d, 1);
        assert_eq!(b"b", Buf::bytes(&d));
        Buf::advance(&mut d, 3);
        assert_eq!(b"e", Buf::bytes(&d));
        Buf::advance(&mut d, 1);
        assert_eq!(0, Buf::remaining(&d));
        assert_eq!(false, Buf::has_remaining(&d));
    }

    #[test]
    fn buf_bytes_vectored() {
        let mut d = BytesVecDeque::new();
        d.extend(Bytes::from_static(b"ab"));
        d.extend(Bytes::from_static(b"cde"));

        let mut v = [IoSlice::new(&[])];
        assert_eq!(1, d.bytes_vectored(&mut v));
        assert_eq!(b"ab", &*v[0]);

        let mut v = [IoSlice::new(&[]), IoSlice::new(&[])];
        assert_eq!(2, d.bytes_vectored(&mut v));
        assert_eq!(b"ab", &*v[0]);
        assert_eq!(b"cde", &*v[1]);

        let mut v = [IoSlice::new(&[]), IoSlice::new(&[]), IoSlice::new(&[])];
        assert_eq!(2, d.bytes_vectored(&mut v));
        assert_eq!(b"ab", &*v[0]);
        assert_eq!(b"cde", &*v[1]);
    }
}
<|endoftext|>


<commit_before>//! Endpoints for user presence.
<commit_msg>Add presence endpoints.<commit_after>//! Endpoints for user presence.

/// PUT /_matrix/client/r0/presence/{userId}/status
///
/// [Matrix spec link](http://matrix.org/docs/spec/client_server/r0.2.0.html#put-matrix-client-r0-presence-userid-status)
pub mod set_presence {
    use ruma_identifiers::UserId;
    use ruma_events::presence::PresenceState;

    /// Details about this API endpoint.
    pub struct Endpoint;

    /// This API endpoint's path parameters.
    #[derive(Clone, Debug, Deserialize, Serialize)]
    pub struct PathParams {
        pub user_id: UserId
    }

    /// This API endpoint's body parameters.
    #[derive(Clone, Debug, Deserialize, Serialize)]
    pub struct BodyParams {
        status_msg: Option<String>,
        presence: PresenceState
    }

    impl ::Endpoint for Endpoint {
        type BodyParams = BodyParams;
        type PathParams = PathParams;
        type QueryParams = ();
        type Response = ();

        fn method() -> ::Method {
            ::Method::Put
        }

        fn request_path(params: Self::PathParams) -> String {
            format!(
                "/_matrix/client/r0/presence/{}/status",
                params.user_id
            )
        }

        fn router_path() -> String {
            "/_matrix/client/r0/presence/:user_id/status".to_string()
        }
    }
}

/// GET /_matrix/client/r0/presence/{userId}/status
///
/// [Matrix spec link](http://matrix.org/docs/spec/client_server/r0.2.0.html#get-matrix-client-r0-presence-userid-status)
pub mod get_presence {
    use ruma_identifiers::UserId;
    use ruma_events::presence::PresenceState;

    /// Details about this API endpoint.
    pub struct Endpoint;

    /// This API endpoint's path parameters.
    #[derive(Clone, Debug, Deserialize, Serialize)]
    pub struct PathParams {
        pub user_id: UserId
    }

    /// This API endpoint's response.
    #[derive(Clone, Debug, Deserialize, Serialize)]
    pub struct Response {
        pub status_msg: Option<String>,
        pub currently_active: Option<bool>,
        pub last_active_ago: Option<u64>,
        pub presence: PresenceState
    }

    impl ::Endpoint for Endpoint {
        type BodyParams = ();
        type PathParams = PathParams;
        type QueryParams = ();
        type Response = Response;

        fn method() -> ::Method {
            ::Method::Get
        }

        fn request_path(params: Self::PathParams) -> String {
            format!(
                "/_matrix/client/r0/presence/{}/status",
                params.user_id
            )
        }

        fn router_path() -> String {
            "/_matrix/client/r0/presence/:user_id/status".to_string()
        }
    }
}

/// POST /_matrix/client/r0/presence/list/{userId}
///
/// [Matrix spec link](http://matrix.org/docs/spec/client_server/r0.2.0.html#post-matrix-client-r0-presence-list-userid)
pub mod update_presence_list {
    use ruma_identifiers::UserId;

    /// Details about this API endpoint.
    pub struct Endpoint;

    /// This API endpoint's path parameters.
    #[derive(Clone, Debug, Deserialize, Serialize)]
    pub struct PathParams {
        pub user_id: UserId
    }

    /// This API endpoint's body parameters.
    #[derive(Clone, Debug, Deserialize, Serialize)]
    pub struct BodyParams {
        drop: Option<Vec<UserId>>,
        invite: Option<Vec<UserId>>
    }

    impl ::Endpoint for Endpoint {
        type BodyParams = BodyParams;
        type PathParams = PathParams;
        type QueryParams = ();
        type Response = ();

        fn method() -> ::Method {
            ::Method::Post
        }

        fn request_path(params: Self::PathParams) -> String {
            format!(
                "/_matrix/client/r0/presence/list/{}",
                params.user_id
            )
        }

        fn router_path() -> String {
            "/_matrix/client/r0/presence/list/:user_id".to_string()
        }
    }
}

/// GET /_matrix/client/r0/presence/list/{userId}
///
/// [Matrix spec link](http://matrix.org/docs/spec/client_server/r0.2.0.html#get-matrix-client-r0-presence-list-userid)
pub mod get_presence_list_status {
    use ruma_identifiers::UserId;
    use ruma_events::presence::PresenceEvent;

    /// Details about this API endpoint.
    pub struct Endpoint;

    /// This API endpoint's path parameters.
    #[derive(Clone, Debug, Deserialize, Serialize)]
    pub struct PathParams {
        pub user_id: UserId
    }

    impl ::Endpoint for Endpoint {
        type BodyParams = ();
        type PathParams = PathParams;
        type QueryParams = ();
        type Response = Vec<PresenceEvent>;

        fn method() -> ::Method {
            ::Method::Get
        }

        fn request_path(params: Self::PathParams) -> String {
            format!(
                "/_matrix/client/r0/presence/list/{}",
                params.user_id
            )
        }

        fn router_path() -> String {
            "/_matrix/client/r0/presence/list/:user_id".to_string()
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

//! This file contains logic for checking a node, common to both VFNs and PFNs.
//! At this point, some earlier code has processed the input information, e.g.
//! VFN information on chain or PFN information from a file, and has converted
//! it into a common format that these functions can ingest.

use aptos_node_checker_lib::EvaluationSummary;
use aptos_sdk::{
    crypto::{x25519, ValidCryptoMaterialStringExt},
    types::account_address::AccountAddress,
};
use clap::Parser;
use futures::stream::FuturesUnordered;
use futures::StreamExt;
use log::{debug, info};
use reqwest::{Client as ReqwestClient, Url};
use serde::Serialize;
use std::collections::HashMap;
use std::{
    sync::Arc,
    time::{Duration, SystemTime, UNIX_EPOCH},
};
use tokio::sync::Semaphore;

// Unfortunately we don't have any way, on chain or not, to know what the API
// port is, so we just assume it is one of these. If their node is inaccessible at,
// any of these ports, e.g. because it runs on 7777 and they have an LB not registered
// on chain in front of it that listens at 80 just for the API, we're just out of luck.
// If we get this kind of information from elsewhere at some point, we could look at
// joining it here.
pub const API_PORTS: &[u16] = &[80, 8080, 443];

#[derive(Debug, Parser)]
pub struct NodeHealthCheckerArgs {
    /// Address where we can expect to find an NHC instance.
    #[clap(long)]
    pub nhc_address: Url,

    /// Baseline config to use when talking to NHC.
    #[clap(long)]
    pub nhc_baseline_config_name: String,

    /// How long to wait when talking to NHC. The check should be quick unless
    /// we're using the TPS evaluator.
    #[clap(long, default_value_t = 60)]
    pub nhc_timeout_secs: u64,

    /// Max number of requests to NHC we can have running concurrently.
    #[clap(long, default_value_t = 64)]
    pub max_concurrent_nhc_requests: u16,
}

impl NodeHealthCheckerArgs {
    /// Check all FullNodes from the validator set.
    pub async fn check_nodes(
        &self,
        address_to_nodes: HashMap<AccountAddress, Vec<NodeInfo>>,
    ) -> HashMap<AccountAddress, Vec<SingleCheck>> {
        let nhc_client = ReqwestClient::builder()
            .timeout(Duration::from_secs(self.nhc_timeout_secs))
            .build()
            .expect("Somehow failed to build reqwest client");

        let mut nhc_address = self.nhc_address.clone();
        nhc_address.set_path("/check_node");

        let semaphore = Arc::new(Semaphore::new(self.max_concurrent_nhc_requests as usize));

        // Build up futures for checking each node.
        let mut nhc_responses = HashMap::new();
        let mut futures = FuturesUnordered::new();
        for (account_address, node_infos) in address_to_nodes {
            for node_info in node_infos {
                let single_check_future = self.check_single_fn_wrapper(
                    account_address,
                    &nhc_client,
                    &nhc_address,
                    node_info,
                    semaphore.clone(),
                );
                futures.push(single_check_future);
            }
        }

        // Go through all the futures, log some information about their results,
        // and insert them into the output.
        while let Some((account_address, node_url, single_check_result)) = futures.next().await {
            match single_check_result {
                SingleCheckResult::Success(_) => {
                    info!("NHC returned a 200 for {}", node_url);
                }
                SingleCheckResult::NodeCheckFailure(_) => {
                    info!("NHC did not return a 200 for {}", node_url);
                }
                wildcard => {
                    panic!(
                        "Shouldn't be possible for checK_single_fn_wrapper to return {:?}",
                        wildcard
                    )
                }
            }
            nhc_responses
                .entry(account_address)
                .or_insert_with(Vec::new)
                .push(SingleCheck::new(single_check_result, Some(node_url)));
        }
        nhc_responses
    }

    /// This just exists to make joining futures easy. This function returns the
    /// account address and node address alongside the check so we can put it in
    /// the output easily.
    async fn check_single_fn_wrapper(
        &self,
        account_address: AccountAddress,
        nhc_client: &ReqwestClient,
        nhc_address: &Url,
        node_info: NodeInfo,
        semaphore: Arc<Semaphore>,
    ) -> (AccountAddress, Url, SingleCheckResult) {
        let _permit = semaphore.acquire().await.unwrap();
        (
            account_address,
            node_info.node_url.clone(),
            self.check_single_fn(nhc_client, nhc_address, node_info)
                .await,
        )
    }

    /// Make a query to NHC for a single node. This may result in multiple
    /// queries to NHC as we try different API ports.
    async fn check_single_fn(
        &self,
        nhc_client: &ReqwestClient,
        nhc_address: &Url,
        node_info: NodeInfo,
    ) -> SingleCheckResult {
        match node_info.api_port {
            Some(api_port) => {
                self.check_single_fn_one_api_port(
                    nhc_client,
                    nhc_address,
                    &node_info.node_url,
                    api_port,
                    node_info.public_key,
                )
                .await
            }
            None => {
                let mut index = 0;
                loop {
                    let single_check_result = self
                        .check_single_fn_one_api_port(
                            nhc_client,
                            nhc_address,
                            &node_info.node_url,
                            API_PORTS[index],
                            node_info.public_key,
                        )
                        .await;
                    // If the response was a success, return.
                    if let SingleCheckResult::Success(_) = &single_check_result {
                        break single_check_result;
                    }
                    // If the response was a failure, return unless it was because it
                    // seems like the API wasn't open, in which case we keep looping
                    // to try another port.
                    if let SingleCheckResult::NodeCheckFailure(failure) = &single_check_result {
                        if failure.code != NodeCheckFailureCode::ApiPortClosed {
                            break single_check_result;
                        }
                    }
                    // We've tried every port, return the last result.
                    if index == API_PORTS.len() - 1 {
                        break single_check_result;
                    }
                    index += 1;
                }
            }
        }
    }

    /// Check NHC just once, with a single node for a single API port.
    async fn check_single_fn_one_api_port(
        &self,
        nhc_client: &ReqwestClient,
        nhc_address: &Url,
        node_url: &Url,
        api_port: u16,
        public_key: Option<x25519::PublicKey>,
    ) -> SingleCheckResult {
        // Build up query params.
        let mut params = HashMap::new();
        params.insert("node_url", node_url.to_string());
        params.insert("api_port", api_port.to_string());
        params.insert(
            "baseline_configuration_name",
            self.nhc_baseline_config_name.clone(),
        );
        if let Some(public_key) = public_key {
            params.insert("public_key", public_key.to_encoded_string().unwrap());
        }

        // This is just for pretty logging / output purposes.
        let address_single_string = {
            let mut node_url = node_url.clone();
            let _ = node_url.set_port(Some(api_port));
            node_url.to_string()
        };
        debug!("Querying NHC at address: {}", address_single_string);

        // Send the request and parse the response.
        let response = match nhc_client
            .get(nhc_address.clone())
            .query(&params)
            .send()
            .await
        {
            Ok(response) => response,
            Err(e) => {
                return SingleCheckResult::NodeCheckFailure(NodeCheckFailure::new(
                    format!("Error with request flow to NHC: {:#}", e),
                    NodeCheckFailureCode::RequestResponseError,
                ));
            }
        };

        // Handle the case where NHC itself throws an error (as opposed to a success
        // response from NHC indicating evaluation where the node performed poorly).
        if let Err(e) = response.error_for_status_ref() {
            return SingleCheckResult::NodeCheckFailure(NodeCheckFailure::new(
                format!("{:#}: {:?}", e, response.text().await),
                NodeCheckFailureCode::ResponseNot200,
            ));
        };

        // Confirm the response is valid JSON.
        let evaluation_summary = match response.json::<EvaluationSummary>().await {
            Ok(evaluation_summary) => evaluation_summary,
            Err(e) => {
                return SingleCheckResult::NodeCheckFailure(NodeCheckFailure::new(
                    format!("{:#}", e),
                    NodeCheckFailureCode::CouldNotDeserializeResponse,
                ))
            }
        };

        // Check specifically if the API port is closed.
        if evaluation_summary
            .evaluation_results
            .iter()
            .any(|evaluation_result| {
                evaluation_result.evaluator_name == "index_response" && evaluation_result.score == 0
            })
        {
            return SingleCheckResult::NodeCheckFailure(NodeCheckFailure::new(
                format!(
                    "Couldn't talk to API on any of the expected ports {:?}: {:?}",
                    API_PORTS, evaluation_summary
                ),
                NodeCheckFailureCode::ApiPortClosed,
            ));
        }

        SingleCheckResult::Success(SingleCheckSuccess::new(
            evaluation_summary,
            address_single_string,
        ))
    }
}

/// Helper struct for putting all the FN address information in one place.
#[derive(Debug)]
pub struct NodeInfo {
    /// This should include the scheme, e.g. http://
    pub node_url: Url,

    /// If given, we will use this. If not, we'll try each of API_PORTS.
    pub api_port: Option<u16>,

    /// If this is included, we'll include this in the NHC request.
    pub public_key: Option<x25519::PublicKey>,
}

#[derive(Debug, Serialize)]
pub struct SingleCheck {
    pub result: SingleCheckResult,
    pub timestamp: Duration,
    pub node_url: Option<Url>,
}

impl SingleCheck {
    pub fn new(result: SingleCheckResult, node_url: Option<Url>) -> Self {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("Failed to get current time");
        Self {
            result,
            timestamp: now,
            node_url,
        }
    }
}

/// We use this struct to capture the result of checking a node, or the lack thereof.
#[derive(Debug, Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum SingleCheckResult {
    /// The node was successfully checked. Note: The evaulation itself could
    /// still indicat, a problem with the node, this just states that we were
    /// able to check the node sucessfully with NHC.
    Success(SingleCheckSuccess),

    /// Something went wrong with checking the node with NHC.
    NodeCheckFailure(NodeCheckFailure),

    /// The account does not have a VFN registered on chain.
    NoVfnRegistered(NoVfnRegistered),

    /// The network address could not be deserialized.
    CouldNotDeserializeNetworkAddress(CouldNotDeserializeNetworkAddress),

    /// The network address was incomplete, e.g. missing an IP, port, public
    /// key, etc.
    IncompleteNetworkAddress(IncompleteNetworkAddress),
}

#[derive(Debug, Serialize)]
pub struct SingleCheckSuccess {
    /// The evaluation summary returned by NHC. This doesn't necessarily imply
    /// that the node passed the evaluation, just that an evaluation was returned
    /// successfully and it passed the API available check.
    pub evaluation_summary: EvaluationSummary,

    /// This is the address that we used to get this successful evaluation.
    /// This is presented in a normal URL format, not the NetworkAddress
    /// representation. Example value for this field: http://65.109.17.29:8080.
    /// Note, sometimes the address we started with was a DNS name, and we resolved
    /// it to an IP address. As such, this IP address may become incorrect down
    /// the line. In that case, refer to fn_address in SingleCheck, or just
    /// run this tool again.
    pub fn_address_url: String,
}

impl SingleCheckSuccess {
    pub fn new(evaluation_summary: EvaluationSummary, fn_address_url: String) -> Self {
        Self {
            evaluation_summary,
            fn_address_url,
        }
    }
}

#[derive(Debug, Serialize)]
pub struct NodeCheckFailure {
    pub message: String,
    pub code: NodeCheckFailureCode,
}

impl NodeCheckFailure {
    pub fn new(message: String, code: NodeCheckFailureCode) -> Self {
        Self { message, code }
    }
}

#[derive(Debug, PartialEq, Eq, Serialize)]
pub enum NodeCheckFailureCode {
    // Something went wrong when sending / receiving the request.
    RequestResponseError,

    // The response from NHC was not a 200, implying a problem with NHC.
    ResponseNot200,

    // The response from NHC couldn't be deserialized.
    CouldNotDeserializeResponse,

    // NHC returned an evaluation that indicates that the API port is closed.
    ApiPortClosed,
}

// These are necessary because we can't just use a unit type for this enum variant.

#[derive(Debug, Serialize)]
pub struct NoVfnRegistered;

#[derive(Debug, Serialize)]
pub struct CouldNotDeserializeNetworkAddress {
    pub message: String,
}

#[derive(Debug, Serialize)]
pub struct IncompleteNetworkAddress {
    pub message: String,
}
<commit_msg>Address nits from #3904 in fn-check-client<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

//! This file contains logic for checking a node, common to both VFNs and PFNs.
//! At this point, some earlier code has processed the input information, e.g.
//! VFN information on chain or PFN information from a file, and has converted
//! it into a common format that these functions can ingest.

use aptos_node_checker_lib::EvaluationSummary;
use aptos_sdk::{
    crypto::{x25519, ValidCryptoMaterialStringExt},
    types::account_address::AccountAddress,
};
use clap::Parser;
use futures::stream::FuturesUnordered;
use futures::StreamExt;
use log::{debug, info};
use reqwest::{Client as ReqwestClient, Url};
use serde::Serialize;
use std::collections::HashMap;
use std::{
    sync::Arc,
    time::{Duration, SystemTime, UNIX_EPOCH},
};
use tokio::sync::Semaphore;

// Unfortunately we don't have any way, on chain or not, to know what the API
// port is, so we just assume it is one of these. If their node is inaccessible at
// any of these ports, e.g. because it runs on 7777 and they have an LB not registered
// on chain in front of it that listens at 80 just for the API, we're just out of luck.
// If we get this kind of information from elsewhere at some point, we could look at
// joining it here.
pub const API_PORTS: &[u16] = &[80, 8080, 443];

#[derive(Debug, Parser)]
pub struct NodeHealthCheckerArgs {
    /// Address where we can expect to find an NHC instance.
    #[clap(long)]
    pub nhc_address: Url,

    /// Baseline config to use when talking to NHC.
    #[clap(long)]
    pub nhc_baseline_config_name: String,

    /// How long to wait when talking to NHC. The check should be quick unless
    /// we're using the TPS evaluator.
    #[clap(long, default_value_t = 60)]
    pub nhc_timeout_secs: u64,

    /// Max number of requests to NHC we can have running concurrently.
    #[clap(long, default_value_t = 64)]
    pub max_concurrent_nhc_requests: u16,
}

impl NodeHealthCheckerArgs {
    /// Per account address, check all nodes. This may result in multiple calls
    /// to NHC even for a single node in the case that we don't know the API
    /// port and we're just guessing.
    pub async fn check_nodes(
        &self,
        address_to_nodes: HashMap<AccountAddress, Vec<NodeInfo>>,
    ) -> HashMap<AccountAddress, Vec<SingleCheck>> {
        let nhc_client = ReqwestClient::builder()
            .timeout(Duration::from_secs(self.nhc_timeout_secs))
            .build()
            .expect("Somehow failed to build reqwest client");

        let mut nhc_address = self.nhc_address.clone();
        nhc_address.set_path("/check_node");

        let semaphore = Arc::new(Semaphore::new(self.max_concurrent_nhc_requests as usize));

        // Build up futures for checking each node.
        let mut nhc_responses = HashMap::new();
        let mut futures = FuturesUnordered::new();
        for (account_address, node_infos) in address_to_nodes {
            for node_info in node_infos {
                let single_check_future = self.check_single_fn_wrapper(
                    account_address,
                    &nhc_client,
                    &nhc_address,
                    node_info,
                    semaphore.clone(),
                );
                futures.push(single_check_future);
            }
        }

        // Go through all the futures, log some information about their results,
        // and insert them into the output.
        while let Some((account_address, node_url, single_check_result)) = futures.next().await {
            match single_check_result {
                SingleCheckResult::Success(_) => {
                    info!("NHC returned a 200 for {}", node_url);
                }
                SingleCheckResult::NodeCheckFailure(_) => {
                    info!("NHC did not return a 200 for {}", node_url);
                }
                wildcard => {
                    panic!(
                        "Shouldn't be possible for checK_single_fn_wrapper to return {:?}",
                        wildcard
                    )
                }
            }
            nhc_responses
                .entry(account_address)
                .or_insert_with(Vec::new)
                .push(SingleCheck::new(single_check_result, Some(node_url)));
        }
        nhc_responses
    }

    /// This just exists to make joining futures easy. This function returns the
    /// account address and node address alongside the check so we can put it in
    /// the output easily.
    async fn check_single_fn_wrapper(
        &self,
        account_address: AccountAddress,
        nhc_client: &ReqwestClient,
        nhc_address: &Url,
        node_info: NodeInfo,
        semaphore: Arc<Semaphore>,
    ) -> (AccountAddress, Url, SingleCheckResult) {
        let _permit = semaphore.acquire().await.unwrap();
        (
            account_address,
            node_info.node_url.clone(),
            self.check_single_fn(nhc_client, nhc_address, node_info)
                .await,
        )
    }

    /// Make a query to NHC for a single node. This may result in multiple
    /// queries to NHC as we try different API ports.
    async fn check_single_fn(
        &self,
        nhc_client: &ReqwestClient,
        nhc_address: &Url,
        node_info: NodeInfo,
    ) -> SingleCheckResult {
        match node_info.api_port {
            Some(api_port) => {
                self.check_single_fn_one_api_port(
                    nhc_client,
                    nhc_address,
                    &node_info.node_url,
                    api_port,
                    node_info.public_key,
                )
                .await
            }
            None => {
                let mut index = 0;
                loop {
                    let single_check_result = self
                        .check_single_fn_one_api_port(
                            nhc_client,
                            nhc_address,
                            &node_info.node_url,
                            API_PORTS[index],
                            node_info.public_key,
                        )
                        .await;
                    // If the response was a success, return.
                    if let SingleCheckResult::Success(_) = &single_check_result {
                        break single_check_result;
                    }
                    // If the response was a failure, return unless it was because it
                    // seems like the API wasn't open, in which case we keep looping
                    // to try another port.
                    if let SingleCheckResult::NodeCheckFailure(failure) = &single_check_result {
                        if failure.code != NodeCheckFailureCode::ApiPortClosed {
                            break single_check_result;
                        }
                    }
                    // We've tried every port, return the last result.
                    if index == API_PORTS.len() - 1 {
                        break single_check_result;
                    }
                    index += 1;
                }
            }
        }
    }

    /// Check NHC just once, with a single node for a single API port.
    async fn check_single_fn_one_api_port(
        &self,
        nhc_client: &ReqwestClient,
        nhc_address: &Url,
        node_url: &Url,
        api_port: u16,
        public_key: Option<x25519::PublicKey>,
    ) -> SingleCheckResult {
        // Build up query params.
        let mut params = HashMap::new();
        params.insert("node_url", node_url.to_string());
        params.insert("api_port", api_port.to_string());
        params.insert(
            "baseline_configuration_name",
            self.nhc_baseline_config_name.clone(),
        );
        if let Some(public_key) = public_key {
            params.insert("public_key", public_key.to_encoded_string().unwrap());
        }

        // This is just for pretty logging / output purposes.
        let address_single_string = {
            let mut node_url = node_url.clone();
            let _ = node_url.set_port(Some(api_port));
            node_url.to_string()
        };
        debug!("Querying NHC at address: {}", address_single_string);

        // Send the request and parse the response.
        let response = match nhc_client
            .get(nhc_address.clone())
            .query(&params)
            .send()
            .await
        {
            Ok(response) => response,
            Err(e) => {
                return SingleCheckResult::NodeCheckFailure(NodeCheckFailure::new(
                    format!("Error with request flow to NHC: {:#}", e),
                    NodeCheckFailureCode::RequestResponseError,
                ));
            }
        };

        // Handle the case where NHC itself throws an error (as opposed to a success
        // response from NHC indicating evaluation where the node performed poorly).
        if let Err(e) = response.error_for_status_ref() {
            return SingleCheckResult::NodeCheckFailure(NodeCheckFailure::new(
                format!("{:#}: {:?}", e, response.text().await),
                NodeCheckFailureCode::ResponseNot200,
            ));
        };

        // Confirm the response is valid JSON.
        let evaluation_summary = match response.json::<EvaluationSummary>().await {
            Ok(evaluation_summary) => evaluation_summary,
            Err(e) => {
                return SingleCheckResult::NodeCheckFailure(NodeCheckFailure::new(
                    format!("{:#}", e),
                    NodeCheckFailureCode::CouldNotDeserializeResponse,
                ))
            }
        };

        // Check specifically if the API port is closed.
        if evaluation_summary
            .evaluation_results
            .iter()
            .any(|evaluation_result| {
                evaluation_result.evaluator_name == "index_response" && evaluation_result.score == 0
            })
        {
            return SingleCheckResult::NodeCheckFailure(NodeCheckFailure::new(
                format!(
                    "Couldn't talk to API on any of the expected ports {:?}: {:?}",
                    API_PORTS, evaluation_summary
                ),
                NodeCheckFailureCode::ApiPortClosed,
            ));
        }

        SingleCheckResult::Success(SingleCheckSuccess::new(
            evaluation_summary,
            address_single_string,
        ))
    }
}

/// Helper struct for putting all the FN address information in one place.
#[derive(Debug)]
pub struct NodeInfo {
    /// This should include the scheme, e.g. http://
    pub node_url: Url,

    /// If given, we will use this. If not, we'll try each of API_PORTS.
    pub api_port: Option<u16>,

    /// If this is included, we'll include this in the NHC request.
    pub public_key: Option<x25519::PublicKey>,
}

#[derive(Debug, Serialize)]
pub struct SingleCheck {
    pub result: SingleCheckResult,
    pub timestamp: Duration,
    pub node_url: Option<Url>,
}

impl SingleCheck {
    pub fn new(result: SingleCheckResult, node_url: Option<Url>) -> Self {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("Failed to get current time");
        Self {
            result,
            timestamp: now,
            node_url,
        }
    }
}

/// We use this struct to capture the result of checking a node, or the lack thereof.
#[derive(Debug, Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum SingleCheckResult {
    /// The node was successfully checked. Note: The evaulation itself could
    /// still indicate a problem with the node, this just states that we were
    /// able to check the node sucessfully with NHC.
    Success(SingleCheckSuccess),

    /// Something went wrong with checking the node with NHC.
    NodeCheckFailure(NodeCheckFailure),

    /// The account does not have a VFN registered on chain.
    NoVfnRegistered(NoVfnRegistered),

    /// The network address could not be deserialized.
    CouldNotDeserializeNetworkAddress(CouldNotDeserializeNetworkAddress),

    /// The network address was incomplete, e.g. missing an IP, port, public
    /// key, etc.
    IncompleteNetworkAddress(IncompleteNetworkAddress),
}

#[derive(Debug, Serialize)]
pub struct SingleCheckSuccess {
    /// The evaluation summary returned by NHC. This doesn't necessarily imply
    /// that the node passed the evaluation, just that an evaluation was returned
    /// successfully and it passed the API available check.
    pub evaluation_summary: EvaluationSummary,

    /// This is the address that we used to get this successful evaluation.
    /// This is presented in a normal URL format, not the NetworkAddress
    /// representation. Example value for this field: http://65.109.17.29:8080.
    /// Note, sometimes the address we started with was a DNS name, and we resolved
    /// it to an IP address. As such, this IP address may become incorrect down
    /// the line. In that case, refer to fn_address in SingleCheck, or just
    /// run this tool again.
    pub fn_address_url: String,
}

impl SingleCheckSuccess {
    pub fn new(evaluation_summary: EvaluationSummary, fn_address_url: String) -> Self {
        Self {
            evaluation_summary,
            fn_address_url,
        }
    }
}

#[derive(Debug, Serialize)]
pub struct NodeCheckFailure {
    pub message: String,
    pub code: NodeCheckFailureCode,
}

impl NodeCheckFailure {
    pub fn new(message: String, code: NodeCheckFailureCode) -> Self {
        Self { message, code }
    }
}

#[derive(Debug, PartialEq, Eq, Serialize)]
pub enum NodeCheckFailureCode {
    // Something went wrong when sending / receiving the request.
    RequestResponseError,

    // The response from NHC was not a 200, implying a problem with NHC.
    ResponseNot200,

    // The response from NHC couldn't be deserialized.
    CouldNotDeserializeResponse,

    // NHC returned an evaluation that indicates that the API port is closed.
    ApiPortClosed,
}

// These are necessary because we can't just use a unit type for this enum variant
// because we serialize SingleCheckResult using internal tagging, in which case
// serde requires that all variants have values.

#[derive(Debug, Serialize)]
pub struct NoVfnRegistered;

#[derive(Debug, Serialize)]
pub struct CouldNotDeserializeNetworkAddress {
    pub message: String,
}

#[derive(Debug, Serialize)]
pub struct IncompleteNetworkAddress {
    pub message: String,
}
<|endoftext|>