<commit_before>// Copyright 2019 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 super::device_slot::{DeviceSlot, DeviceSlots, Error as DeviceSlotError};
use super::interrupter::{Error as InterrupterError, Interrupter};
use super::ring_buffer_controller::{
    Error as RingBufferControllerError, RingBufferController, TransferDescriptorHandler,
};
use super::xhci_abi::{
    AddressDeviceCommandTrb, AddressedTrb, ConfigureEndpointCommandTrb, DisableSlotCommandTrb,
    Error as TrbError, EvaluateContextCommandTrb, ResetDeviceCommandTrb,
    SetTRDequeuePointerCommandTrb, StopEndpointCommandTrb, TransferDescriptor, TrbCast,
    TrbCompletionCode, TrbType,
};
use super::xhci_regs::{valid_slot_id, MAX_SLOTS};
use std::fmt::{self, Display};
use std::sync::Arc;
use sync::Mutex;
use sys_util::{Error as SysError, EventFd, GuestAddress, GuestMemory};
use utils::EventLoop;

#[derive(Debug)]
pub enum Error {
    WriteEventFd(SysError),
    SendInterrupt(InterrupterError),
    CastTrb(TrbError),
    BadSlotId(u8),
    StopEndpoint(DeviceSlotError),
    ConfigEndpoint(DeviceSlotError),
    SetAddress(DeviceSlotError),
    EvaluateContext(DeviceSlotError),
    DisableSlot(DeviceSlotError),
    ResetSlot(DeviceSlotError),
}

type Result<T> = std::result::Result<T, Error>;

impl Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Error::*;

        match self {
            WriteEventFd(e) => write!(f, "failed to write event fd: {}", e),
            SendInterrupt(e) => write!(f, "failed to send interrupt: {}", e),
            CastTrb(e) => write!(f, "failed to cast trb: {}", e),
            BadSlotId(id) => write!(f, "bad slot id: {}", id),
            StopEndpoint(e) => write!(f, "failed to stop endpoint: {}", e),
            ConfigEndpoint(e) => write!(f, "failed to config endpoint: {}", e),
            SetAddress(e) => write!(f, "failed to set address: {}", e),
            EvaluateContext(e) => write!(f, "failed to evaluate context: {}", e),
            DisableSlot(e) => write!(f, "failed to disable slot: {}", e),
            ResetSlot(e) => write!(f, "failed to reset slot: {}", e),
        }
    }
}

pub type CommandRingController = RingBufferController<CommandRingTrbHandler>;
pub type CommandRingControllerError = RingBufferControllerError;

impl CommandRingController {
    pub fn new(
        mem: GuestMemory,
        event_loop: Arc<EventLoop>,
        slots: DeviceSlots,
        interrupter: Arc<Mutex<Interrupter>>,
    ) -> std::result::Result<Arc<CommandRingController>, RingBufferControllerError> {
        RingBufferController::new_with_handler(
            String::from("command ring"),
            mem,
            event_loop,
            CommandRingTrbHandler::new(slots, interrupter),
        )
    }
}

pub struct CommandRingTrbHandler {
    slots: DeviceSlots,
    interrupter: Arc<Mutex<Interrupter>>,
}

impl CommandRingTrbHandler {
    fn new(slots: DeviceSlots, interrupter: Arc<Mutex<Interrupter>>) -> Self {
        CommandRingTrbHandler { slots, interrupter }
    }

    fn slot(&self, slot_id: u8) -> Result<Arc<DeviceSlot>> {
        self.slots.slot(slot_id).ok_or(Error::BadSlotId(slot_id))
    }

    fn command_completion_callback(
        interrupter: &Arc<Mutex<Interrupter>>,
        completion_code: TrbCompletionCode,
        slot_id: u8,
        trb_addr: u64,
        event_fd: &EventFd,
    ) -> Result<()> {
        interrupter
            .lock()
            .send_command_completion_trb(completion_code, slot_id, GuestAddress(trb_addr))
            .map_err(Error::SendInterrupt)?;
        event_fd.write(1).map_err(Error::WriteEventFd)
    }

    fn enable_slot(&self, atrb: &AddressedTrb, event_fd: EventFd) -> Result<()> {
        for slot_id in 1..=MAX_SLOTS {
            if self.slot(slot_id)?.enable() {
                return CommandRingTrbHandler::command_completion_callback(
                    &self.interrupter,
                    TrbCompletionCode::Success,
                    slot_id,
                    atrb.gpa,
                    &event_fd,
                );
            }
        }

        CommandRingTrbHandler::command_completion_callback(
            &self.interrupter,
            TrbCompletionCode::NoSlotsAvailableError,
            0,
            atrb.gpa,
            &event_fd,
        )
    }

    fn disable_slot(&self, atrb: &AddressedTrb, event_fd: EventFd) -> Result<()> {
        let trb = atrb
            .trb
            .cast::<DisableSlotCommandTrb>()
            .map_err(Error::CastTrb)?;
        let slot_id = trb.get_slot_id();
        if valid_slot_id(slot_id) {
            let gpa = atrb.gpa;
            let interrupter = self.interrupter.clone();
            self.slots
                .disable_slot(slot_id, move |completion_code| {
                    CommandRingTrbHandler::command_completion_callback(
                        &interrupter,
                        completion_code,
                        slot_id,
                        gpa,
                        &event_fd,
                    )
                    .map_err(|e| {
                        error!("failed to run command completion callback: {}", e);
                        ()
                    })
                })
                .map_err(Error::DisableSlot)
        } else {
            CommandRingTrbHandler::command_completion_callback(
                &self.interrupter,
                TrbCompletionCode::TrbError,
                slot_id,
                atrb.gpa,
                &event_fd,
            )
        }
    }

    fn address_device(&self, atrb: &AddressedTrb, event_fd: EventFd) -> Result<()> {
        let trb = atrb
            .trb
            .cast::<AddressDeviceCommandTrb>()
            .map_err(Error::CastTrb)?;
        let slot_id = trb.get_slot_id();
        let completion_code = {
            if valid_slot_id(slot_id) {
                self.slot(slot_id)?
                    .set_address(trb)
                    .map_err(Error::SetAddress)?
            } else {
                TrbCompletionCode::TrbError
            }
        };
        CommandRingTrbHandler::command_completion_callback(
            &self.interrupter,
            completion_code,
            slot_id,
            atrb.gpa,
            &event_fd,
        )
    }

    fn configure_endpoint(&self, atrb: &AddressedTrb, event_fd: EventFd) -> Result<()> {
        let trb = atrb
            .trb
            .cast::<ConfigureEndpointCommandTrb>()
            .map_err(Error::CastTrb)?;
        let slot_id = trb.get_slot_id();
        let completion_code = {
            if valid_slot_id(slot_id) {
                self.slot(slot_id)?
                    .configure_endpoint(trb)
                    .map_err(Error::ConfigEndpoint)?
            } else {
                TrbCompletionCode::TrbError
            }
        };
        CommandRingTrbHandler::command_completion_callback(
            &self.interrupter,
            completion_code,
            slot_id,
            atrb.gpa,
            &event_fd,
        )
    }

    fn evaluate_context(&self, atrb: &AddressedTrb, event_fd: EventFd) -> Result<()> {
        let trb = atrb
            .trb
            .cast::<EvaluateContextCommandTrb>()
            .map_err(Error::CastTrb)?;
        let slot_id = trb.get_slot_id();
        let completion_code = {
            if valid_slot_id(slot_id) {
                self.slot(slot_id)?
                    .evaluate_context(trb)
                    .map_err(Error::EvaluateContext)?
            } else {
                TrbCompletionCode::TrbError
            }
        };
        CommandRingTrbHandler::command_completion_callback(
            &self.interrupter,
            completion_code,
            slot_id,
            atrb.gpa,
            &event_fd,
        )
    }

    fn reset_device(&self, atrb: &AddressedTrb, event_fd: EventFd) -> Result<()> {
        let trb = atrb
            .trb
            .cast::<ResetDeviceCommandTrb>()
            .map_err(Error::CastTrb)?;
        let slot_id = trb.get_slot_id();
        if valid_slot_id(slot_id) {
            let gpa = atrb.gpa;
            let interrupter = self.interrupter.clone();
            self.slots
                .reset_slot(slot_id, move |completion_code| {
                    CommandRingTrbHandler::command_completion_callback(
                        &interrupter,
                        completion_code,
                        slot_id,
                        gpa,
                        &event_fd,
                    )
                    .map_err(|e| {
                        error!("command completion callback failed: {}", e);
                        ()
                    })
                })
                .map_err(Error::ResetSlot)
        } else {
            CommandRingTrbHandler::command_completion_callback(
                &self.interrupter,
                TrbCompletionCode::TrbError,
                slot_id,
                atrb.gpa,
                &event_fd,
            )
        }
    }

    fn stop_endpoint(&self, atrb: &AddressedTrb, event_fd: EventFd) -> Result<()> {
        let trb = atrb
            .trb
            .cast::<StopEndpointCommandTrb>()
            .map_err(Error::CastTrb)?;
        let slot_id = trb.get_slot_id();
        let endpoint_id = trb.get_endpoint_id();
        if valid_slot_id(slot_id) {
            let gpa = atrb.gpa;
            let interrupter = self.interrupter.clone();
            self.slots
                .stop_endpoint(slot_id, endpoint_id, move |completion_code| {
                    CommandRingTrbHandler::command_completion_callback(
                        &interrupter,
                        completion_code,
                        slot_id,
                        gpa,
                        &event_fd,
                    )
                    .map_err(|e| {
                        error!("command completion callback failed: {}", e);
                        ()
                    })
                })
                .map_err(Error::StopEndpoint)?;
            Ok(())
        } else {
            error!("stop endpoint trb has invalid slot id {}", slot_id);
            CommandRingTrbHandler::command_completion_callback(
                &self.interrupter,
                TrbCompletionCode::TrbError,
                slot_id,
                atrb.gpa,
                &event_fd,
            )
        }
    }

    fn set_tr_dequeue_ptr(&self, atrb: &AddressedTrb, event_fd: EventFd) -> Result<()> {
        let trb = atrb
            .trb
            .cast::<SetTRDequeuePointerCommandTrb>()
            .map_err(Error::CastTrb)?;
        let slot_id = trb.get_slot_id();
        let endpoint_id = trb.get_trb_type();
        // See Set TR Dequeue Pointer Trb in spec.
        let dequeue_ptr = trb.get_dequeue_ptr() << 4;
        let completion_code = {
            if valid_slot_id(slot_id) {
                self.slot(slot_id)?
                    .set_tr_dequeue_ptr(endpoint_id, dequeue_ptr)
            } else {
                error!("stop endpoint trb has invalid slot id {}", slot_id);
                TrbCompletionCode::TrbError
            }
        };
        CommandRingTrbHandler::command_completion_callback(
            &self.interrupter,
            completion_code,
            slot_id,
            atrb.gpa,
            &event_fd,
        )
    }
}

impl TransferDescriptorHandler for CommandRingTrbHandler {
    fn handle_transfer_descriptor(
        &self,
        descriptor: TransferDescriptor,
        complete_event: EventFd,
    ) -> std::result::Result<(), ()> {
        // Command descriptor always consist of a single TRB.
        assert_eq!(descriptor.len(), 1);
        let atrb = &descriptor[0];
        let command_result = match atrb.trb.trb_type() {
            Ok(TrbType::EnableSlotCommand) => self.enable_slot(atrb, complete_event),
            Ok(TrbType::DisableSlotCommand) => self.disable_slot(atrb, complete_event),
            Ok(TrbType::AddressDeviceCommand) => self.address_device(atrb, complete_event),
            Ok(TrbType::ConfigureEndpointCommand) => self.configure_endpoint(atrb, complete_event),
            Ok(TrbType::EvaluateContextCommand) => self.evaluate_context(atrb, complete_event),
            Ok(TrbType::ResetDeviceCommand) => self.reset_device(atrb, complete_event),
            Ok(TrbType::NoopCommand) => CommandRingTrbHandler::command_completion_callback(
                &self.interrupter,
                TrbCompletionCode::Success,
                0,
                atrb.gpa,
                &complete_event,
            ),
            Ok(TrbType::ResetEndpointCommand) => {
                error!(
                    "Receiving reset endpoint command. \
                     It should only happen when cmd ring stall"
                );
                CommandRingTrbHandler::command_completion_callback(
                    &self.interrupter,
                    TrbCompletionCode::TrbError,
                    0,
                    atrb.gpa,
                    &complete_event,
                )
            }
            Ok(TrbType::StopEndpointCommand) => self.stop_endpoint(atrb, complete_event),
            Ok(TrbType::SetTRDequeuePointerCommand) => {
                self.set_tr_dequeue_ptr(atrb, complete_event)
            }
            _ => {
                warn!(
                    // We are not handling type 14,15,16. See table 6.4.6.
                    "Unexpected command ring trb type: {}",
                    atrb.trb.get_trb_type()
                );
                match self.interrupter.lock().send_command_completion_trb(
                    TrbCompletionCode::TrbError,
                    0,
                    GuestAddress(atrb.gpa),
                ) {
                    Err(e) => Err(Error::SendInterrupt(e)),
                    Ok(_) => complete_event.write(1).map_err(Error::WriteEventFd),
                }
            }
        };
        command_result.map_err(|e| {
            error!("command failed: {}", e);
            ()
        })
    }
}
<commit_msg>fix endpoint id in set tr dequeue pointer<commit_after>// Copyright 2019 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 super::device_slot::{DeviceSlot, DeviceSlots, Error as DeviceSlotError};
use super::interrupter::{Error as InterrupterError, Interrupter};
use super::ring_buffer_controller::{
    Error as RingBufferControllerError, RingBufferController, TransferDescriptorHandler,
};
use super::xhci_abi::{
    AddressDeviceCommandTrb, AddressedTrb, ConfigureEndpointCommandTrb, DisableSlotCommandTrb,
    Error as TrbError, EvaluateContextCommandTrb, ResetDeviceCommandTrb,
    SetTRDequeuePointerCommandTrb, StopEndpointCommandTrb, TransferDescriptor, TrbCast,
    TrbCompletionCode, TrbType,
};
use super::xhci_regs::{valid_slot_id, MAX_SLOTS};
use std::fmt::{self, Display};
use std::sync::Arc;
use sync::Mutex;
use sys_util::{Error as SysError, EventFd, GuestAddress, GuestMemory};
use utils::EventLoop;

#[derive(Debug)]
pub enum Error {
    WriteEventFd(SysError),
    SendInterrupt(InterrupterError),
    CastTrb(TrbError),
    BadSlotId(u8),
    StopEndpoint(DeviceSlotError),
    ConfigEndpoint(DeviceSlotError),
    SetAddress(DeviceSlotError),
    EvaluateContext(DeviceSlotError),
    DisableSlot(DeviceSlotError),
    ResetSlot(DeviceSlotError),
}

type Result<T> = std::result::Result<T, Error>;

impl Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Error::*;

        match self {
            WriteEventFd(e) => write!(f, "failed to write event fd: {}", e),
            SendInterrupt(e) => write!(f, "failed to send interrupt: {}", e),
            CastTrb(e) => write!(f, "failed to cast trb: {}", e),
            BadSlotId(id) => write!(f, "bad slot id: {}", id),
            StopEndpoint(e) => write!(f, "failed to stop endpoint: {}", e),
            ConfigEndpoint(e) => write!(f, "failed to config endpoint: {}", e),
            SetAddress(e) => write!(f, "failed to set address: {}", e),
            EvaluateContext(e) => write!(f, "failed to evaluate context: {}", e),
            DisableSlot(e) => write!(f, "failed to disable slot: {}", e),
            ResetSlot(e) => write!(f, "failed to reset slot: {}", e),
        }
    }
}

pub type CommandRingController = RingBufferController<CommandRingTrbHandler>;
pub type CommandRingControllerError = RingBufferControllerError;

impl CommandRingController {
    pub fn new(
        mem: GuestMemory,
        event_loop: Arc<EventLoop>,
        slots: DeviceSlots,
        interrupter: Arc<Mutex<Interrupter>>,
    ) -> std::result::Result<Arc<CommandRingController>, RingBufferControllerError> {
        RingBufferController::new_with_handler(
            String::from("command ring"),
            mem,
            event_loop,
            CommandRingTrbHandler::new(slots, interrupter),
        )
    }
}

pub struct CommandRingTrbHandler {
    slots: DeviceSlots,
    interrupter: Arc<Mutex<Interrupter>>,
}

impl CommandRingTrbHandler {
    fn new(slots: DeviceSlots, interrupter: Arc<Mutex<Interrupter>>) -> Self {
        CommandRingTrbHandler { slots, interrupter }
    }

    fn slot(&self, slot_id: u8) -> Result<Arc<DeviceSlot>> {
        self.slots.slot(slot_id).ok_or(Error::BadSlotId(slot_id))
    }

    fn command_completion_callback(
        interrupter: &Arc<Mutex<Interrupter>>,
        completion_code: TrbCompletionCode,
        slot_id: u8,
        trb_addr: u64,
        event_fd: &EventFd,
    ) -> Result<()> {
        interrupter
            .lock()
            .send_command_completion_trb(completion_code, slot_id, GuestAddress(trb_addr))
            .map_err(Error::SendInterrupt)?;
        event_fd.write(1).map_err(Error::WriteEventFd)
    }

    fn enable_slot(&self, atrb: &AddressedTrb, event_fd: EventFd) -> Result<()> {
        for slot_id in 1..=MAX_SLOTS {
            if self.slot(slot_id)?.enable() {
                return CommandRingTrbHandler::command_completion_callback(
                    &self.interrupter,
                    TrbCompletionCode::Success,
                    slot_id,
                    atrb.gpa,
                    &event_fd,
                );
            }
        }

        CommandRingTrbHandler::command_completion_callback(
            &self.interrupter,
            TrbCompletionCode::NoSlotsAvailableError,
            0,
            atrb.gpa,
            &event_fd,
        )
    }

    fn disable_slot(&self, atrb: &AddressedTrb, event_fd: EventFd) -> Result<()> {
        let trb = atrb
            .trb
            .cast::<DisableSlotCommandTrb>()
            .map_err(Error::CastTrb)?;
        let slot_id = trb.get_slot_id();
        if valid_slot_id(slot_id) {
            let gpa = atrb.gpa;
            let interrupter = self.interrupter.clone();
            self.slots
                .disable_slot(slot_id, move |completion_code| {
                    CommandRingTrbHandler::command_completion_callback(
                        &interrupter,
                        completion_code,
                        slot_id,
                        gpa,
                        &event_fd,
                    )
                    .map_err(|e| {
                        error!("failed to run command completion callback: {}", e);
                        ()
                    })
                })
                .map_err(Error::DisableSlot)
        } else {
            CommandRingTrbHandler::command_completion_callback(
                &self.interrupter,
                TrbCompletionCode::TrbError,
                slot_id,
                atrb.gpa,
                &event_fd,
            )
        }
    }

    fn address_device(&self, atrb: &AddressedTrb, event_fd: EventFd) -> Result<()> {
        let trb = atrb
            .trb
            .cast::<AddressDeviceCommandTrb>()
            .map_err(Error::CastTrb)?;
        let slot_id = trb.get_slot_id();
        let completion_code = {
            if valid_slot_id(slot_id) {
                self.slot(slot_id)?
                    .set_address(trb)
                    .map_err(Error::SetAddress)?
            } else {
                TrbCompletionCode::TrbError
            }
        };
        CommandRingTrbHandler::command_completion_callback(
            &self.interrupter,
            completion_code,
            slot_id,
            atrb.gpa,
            &event_fd,
        )
    }

    fn configure_endpoint(&self, atrb: &AddressedTrb, event_fd: EventFd) -> Result<()> {
        let trb = atrb
            .trb
            .cast::<ConfigureEndpointCommandTrb>()
            .map_err(Error::CastTrb)?;
        let slot_id = trb.get_slot_id();
        let completion_code = {
            if valid_slot_id(slot_id) {
                self.slot(slot_id)?
                    .configure_endpoint(trb)
                    .map_err(Error::ConfigEndpoint)?
            } else {
                TrbCompletionCode::TrbError
            }
        };
        CommandRingTrbHandler::command_completion_callback(
            &self.interrupter,
            completion_code,
            slot_id,
            atrb.gpa,
            &event_fd,
        )
    }

    fn evaluate_context(&self, atrb: &AddressedTrb, event_fd: EventFd) -> Result<()> {
        let trb = atrb
            .trb
            .cast::<EvaluateContextCommandTrb>()
            .map_err(Error::CastTrb)?;
        let slot_id = trb.get_slot_id();
        let completion_code = {
            if valid_slot_id(slot_id) {
                self.slot(slot_id)?
                    .evaluate_context(trb)
                    .map_err(Error::EvaluateContext)?
            } else {
                TrbCompletionCode::TrbError
            }
        };
        CommandRingTrbHandler::command_completion_callback(
            &self.interrupter,
            completion_code,
            slot_id,
            atrb.gpa,
            &event_fd,
        )
    }

    fn reset_device(&self, atrb: &AddressedTrb, event_fd: EventFd) -> Result<()> {
        let trb = atrb
            .trb
            .cast::<ResetDeviceCommandTrb>()
            .map_err(Error::CastTrb)?;
        let slot_id = trb.get_slot_id();
        if valid_slot_id(slot_id) {
            let gpa = atrb.gpa;
            let interrupter = self.interrupter.clone();
            self.slots
                .reset_slot(slot_id, move |completion_code| {
                    CommandRingTrbHandler::command_completion_callback(
                        &interrupter,
                        completion_code,
                        slot_id,
                        gpa,
                        &event_fd,
                    )
                    .map_err(|e| {
                        error!("command completion callback failed: {}", e);
                        ()
                    })
                })
                .map_err(Error::ResetSlot)
        } else {
            CommandRingTrbHandler::command_completion_callback(
                &self.interrupter,
                TrbCompletionCode::TrbError,
                slot_id,
                atrb.gpa,
                &event_fd,
            )
        }
    }

    fn stop_endpoint(&self, atrb: &AddressedTrb, event_fd: EventFd) -> Result<()> {
        let trb = atrb
            .trb
            .cast::<StopEndpointCommandTrb>()
            .map_err(Error::CastTrb)?;
        let slot_id = trb.get_slot_id();
        let endpoint_id = trb.get_endpoint_id();
        if valid_slot_id(slot_id) {
            let gpa = atrb.gpa;
            let interrupter = self.interrupter.clone();
            self.slots
                .stop_endpoint(slot_id, endpoint_id, move |completion_code| {
                    CommandRingTrbHandler::command_completion_callback(
                        &interrupter,
                        completion_code,
                        slot_id,
                        gpa,
                        &event_fd,
                    )
                    .map_err(|e| {
                        error!("command completion callback failed: {}", e);
                        ()
                    })
                })
                .map_err(Error::StopEndpoint)?;
            Ok(())
        } else {
            error!("stop endpoint trb has invalid slot id {}", slot_id);
            CommandRingTrbHandler::command_completion_callback(
                &self.interrupter,
                TrbCompletionCode::TrbError,
                slot_id,
                atrb.gpa,
                &event_fd,
            )
        }
    }

    fn set_tr_dequeue_ptr(&self, atrb: &AddressedTrb, event_fd: EventFd) -> Result<()> {
        let trb = atrb
            .trb
            .cast::<SetTRDequeuePointerCommandTrb>()
            .map_err(Error::CastTrb)?;
        let slot_id = trb.get_slot_id();
        let endpoint_id = trb.get_endpoint_id();
        // See Set TR Dequeue Pointer Trb in spec.
        let dequeue_ptr = trb.get_dequeue_ptr() << 4;
        let completion_code = {
            if valid_slot_id(slot_id) {
                self.slot(slot_id)?
                    .set_tr_dequeue_ptr(endpoint_id, dequeue_ptr)
            } else {
                error!("stop endpoint trb has invalid slot id {}", slot_id);
                TrbCompletionCode::TrbError
            }
        };
        CommandRingTrbHandler::command_completion_callback(
            &self.interrupter,
            completion_code,
            slot_id,
            atrb.gpa,
            &event_fd,
        )
    }
}

impl TransferDescriptorHandler for CommandRingTrbHandler {
    fn handle_transfer_descriptor(
        &self,
        descriptor: TransferDescriptor,
        complete_event: EventFd,
    ) -> std::result::Result<(), ()> {
        // Command descriptor always consist of a single TRB.
        assert_eq!(descriptor.len(), 1);
        let atrb = &descriptor[0];
        let command_result = match atrb.trb.trb_type() {
            Ok(TrbType::EnableSlotCommand) => self.enable_slot(atrb, complete_event),
            Ok(TrbType::DisableSlotCommand) => self.disable_slot(atrb, complete_event),
            Ok(TrbType::AddressDeviceCommand) => self.address_device(atrb, complete_event),
            Ok(TrbType::ConfigureEndpointCommand) => self.configure_endpoint(atrb, complete_event),
            Ok(TrbType::EvaluateContextCommand) => self.evaluate_context(atrb, complete_event),
            Ok(TrbType::ResetDeviceCommand) => self.reset_device(atrb, complete_event),
            Ok(TrbType::NoopCommand) => CommandRingTrbHandler::command_completion_callback(
                &self.interrupter,
                TrbCompletionCode::Success,
                0,
                atrb.gpa,
                &complete_event,
            ),
            Ok(TrbType::ResetEndpointCommand) => {
                error!(
                    "Receiving reset endpoint command. \
                     It should only happen when cmd ring stall"
                );
                CommandRingTrbHandler::command_completion_callback(
                    &self.interrupter,
                    TrbCompletionCode::TrbError,
                    0,
                    atrb.gpa,
                    &complete_event,
                )
            }
            Ok(TrbType::StopEndpointCommand) => self.stop_endpoint(atrb, complete_event),
            Ok(TrbType::SetTRDequeuePointerCommand) => {
                self.set_tr_dequeue_ptr(atrb, complete_event)
            }
            _ => {
                warn!(
                    // We are not handling type 14,15,16. See table 6.4.6.
                    "Unexpected command ring trb type: {}",
                    atrb.trb.get_trb_type()
                );
                match self.interrupter.lock().send_command_completion_trb(
                    TrbCompletionCode::TrbError,
                    0,
                    GuestAddress(atrb.gpa),
                ) {
                    Err(e) => Err(Error::SendInterrupt(e)),
                    Ok(_) => complete_event.write(1).map_err(Error::WriteEventFd),
                }
            }
        };
        command_result.map_err(|e| {
            error!("command failed: {}", e);
            ()
        })
    }
}
<|endoftext|>


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

use serde::{Deserialize, Serialize};
use slog::{o, warn, Logger};

use metrics::Metrics;

use crate::extensions::filters::CAPTURED_BYTES;
use crate::extensions::{
    CreateFilterArgs, DownstreamContext, DownstreamResponse, Error, Filter, FilterFactory,
    UpstreamContext, UpstreamResponse,
};

mod metrics;

#[derive(Serialize, Deserialize, Debug)]
/// Strategy to apply for acquiring a set of bytes in the UDP packet
enum Strategy {
    #[serde(rename = "PREFIX")]
    /// Looks for the set of bytes at the beginning of the packet
    Prefix,
    #[serde(rename = "SUFFIX")]
    /// Look for the set of bytes at the end of the packet
    Suffix,
}

#[derive(Serialize, Deserialize, Debug)]
struct Config {
    #[serde(default)]
    strategy: Strategy,
    /// the number of bytes to capture
    #[serde(rename = "size")]
    size: usize,
    /// the key to use when storing the captured bytes in the filter context
    #[serde(rename = "metadataKey")]
    #[serde(default = "default_metadata_key")]
    metadata_key: String,
    /// whether or not to remove the set of the bytes from the packet once captured
    #[serde(default)]
    remove: bool,
}

/// default value for the context key in the Config
fn default_metadata_key() -> String {
    CAPTURED_BYTES.into()
}

impl Default for Strategy {
    fn default() -> Self {
        Strategy::Suffix
    }
}

pub struct CaptureBytesFactory {
    log: Logger,
}

impl CaptureBytesFactory {
    pub fn new(base: &Logger) -> Self {
        CaptureBytesFactory { log: base.clone() }
    }
}

impl FilterFactory for CaptureBytesFactory {
    fn name(&self) -> String {
        "quilkin.extensions.filters.capture_bytes.v1alpha1.CaptureBytes".into()
    }

    fn create_filter(&self, args: CreateFilterArgs) -> Result<Box<dyn Filter>, Error> {
        let config: Config = serde_yaml::to_string(&args.config)
            .and_then(|raw_config| serde_yaml::from_str(raw_config.as_str()))
            .map_err(|err| Error::DeserializeFailed(err.to_string()))?;

        Ok(Box::new(CaptureBytes::new(
            &self.log,
            config,
            Metrics::new(&args.metrics_registry)?,
        )))
    }
}

struct CaptureBytes {
    log: Logger,
    capture: Box<dyn Capture + Sync + Send>,
    /// metrics reporter for this filter.
    metrics: Metrics,
    context_key: String,
    size: usize,
    remove: bool,
}

impl CaptureBytes {
    fn new(base: &Logger, config: Config, metrics: Metrics) -> Self {
        let capture: Box<dyn Capture + Sync + Send> = match config.strategy {
            Strategy::Prefix => Box::new(Prefix {}),
            Strategy::Suffix => Box::new(Suffix {}),
        };

        CaptureBytes {
            log: base.new(o!("source" => "extensions::CaptureBytes")),
            capture,
            metrics,
            context_key: config.metadata_key,
            size: config.size,
            remove: config.remove,
        }
    }
}

impl Filter for CaptureBytes {
    fn on_downstream_receive(&self, mut ctx: DownstreamContext) -> Option<DownstreamResponse> {
        // if the capture size is bigger than the packet size, then we drop the packet,
        // and occasionally warn
        if ctx.contents.len() < self.size {
            if self.metrics.packets_dropped_total.get() % 1000 == 0 {
                warn!(
                    self.log,
                    "packets are being dropped due to their length being less than {} bytes",
                    self.size; "count" => self.metrics.packets_dropped_total.get()
                );
            }
            self.metrics.packets_dropped_total.inc();
            return None;
        }
        let token = self
            .capture
            .capture(&mut ctx.contents, self.size, self.remove);

        ctx.metadata
            .insert(self.context_key.clone(), Box::new(token));

        Some(ctx.into())
    }

    fn on_upstream_receive(&self, ctx: UpstreamContext) -> Option<UpstreamResponse> {
        Some(ctx.into())
    }
}

/// Trait to implement different strategies for capturing packet data
trait Capture {
    /// Capture the packet data from the contents. If remove is true, contents will be altered to
    /// not have the retrieved set of bytes.
    /// Returns the captured bytes.
    fn capture(&self, contents: &mut Vec<u8>, size: usize, remove: bool) -> Vec<u8>;
}

struct Suffix;
impl Capture for Suffix {
    fn capture(&self, contents: &mut Vec<u8>, size: usize, remove: bool) -> Vec<u8> {
        if remove {
            return contents.split_off(contents.len() - size);
        }

        contents
            .iter()
            .skip(contents.len() - size)
            .cloned()
            .collect::<Vec<u8>>()
    }
}

struct Prefix;
impl Capture for Prefix {
    fn capture(&self, contents: &mut Vec<u8>, size: usize, remove: bool) -> Vec<u8> {
        if remove {
            return contents.drain(..size).collect();
        }

        contents.iter().cloned().take(size).collect()
    }
}

#[cfg(test)]
mod tests {
    use prometheus::Registry;
    use serde_yaml::{Mapping, Value};

    use crate::config::{ConnectionConfig, EndPoint};
    use crate::test_utils::{assert_filter_on_downstream_receive_no_change, logger};

    use super::*;

    const TOKEN_KEY: &str = "TOKEN";

    fn capture_bytes(config: Config) -> CaptureBytes {
        CaptureBytes::new(
            &logger(),
            config,
            Metrics::new(&Registry::default()).unwrap(),
        )
    }

    #[test]
    fn factory_valid_config_all() {
        let factory = CaptureBytesFactory::new(&logger());
        let connection = ConnectionConfig::Server { endpoints: vec![] };
        let mut map = Mapping::new();
        map.insert(
            Value::String("strategy".into()),
            Value::String("SUFFIX".into()),
        );
        map.insert(
            Value::String("metadataKey".into()),
            Value::String(TOKEN_KEY.into()),
        );
        map.insert(Value::String("size".into()), Value::Number(3.into()));
        map.insert(Value::String("remove".into()), Value::Bool(true));

        let filter = factory
            .create_filter(CreateFilterArgs::new(
                &connection,
                Some(&Value::Mapping(map)),
            ))
            .unwrap();
        assert_end_strategy(filter.as_ref(), TOKEN_KEY, true);
    }

    #[test]
    fn factory_valid_config_defaults() {
        let factory = CaptureBytesFactory::new(&logger());
        let connection = ConnectionConfig::Server { endpoints: vec![] };
        let mut map = Mapping::new();
        map.insert(Value::String("size".into()), Value::Number(3.into()));
        let filter = factory
            .create_filter(CreateFilterArgs::new(
                &connection,
                Some(&Value::Mapping(map)),
            ))
            .unwrap();
        assert_end_strategy(filter.as_ref(), CAPTURED_BYTES, false);
    }

    #[test]
    fn factory_invalid_config() {
        let factory = CaptureBytesFactory::new(&logger());
        let connection = ConnectionConfig::Server { endpoints: vec![] };
        let mut map = Mapping::new();
        map.insert(Value::String("size".into()), Value::String("WRONG".into()));

        let result = factory.create_filter(CreateFilterArgs::new(
            &connection,
            Some(&Value::Mapping(map)),
        ));
        assert!(result.is_err(), "Should be an error");
    }

    #[test]
    fn on_downstream_receive() {
        let config = Config {
            strategy: Strategy::Suffix,
            metadata_key: TOKEN_KEY.into(),
            size: 3,
            remove: true,
        };
        let filter = capture_bytes(config);
        assert_end_strategy(&filter, TOKEN_KEY, true);
    }

    #[test]
    fn on_downstream_receive_overflow_capture_size() {
        let config = Config {
            strategy: Strategy::Suffix,
            metadata_key: TOKEN_KEY.into(),
            size: 99,
            remove: true,
        };
        let filter = capture_bytes(config);
        let endpoints = vec![EndPoint {
            name: "e1".to_string(),
            address: "127.0.0.1:81".parse().unwrap(),
            connection_ids: vec![],
        }];
        let response = filter.on_downstream_receive(DownstreamContext::new(
            endpoints,
            "127.0.0.1:80".parse().unwrap(),
            "abc".to_string().into_bytes(),
        ));

        assert!(response.is_none());
        let count = filter.metrics.packets_dropped_total.get();
        assert_eq!(1, count);
    }

    #[test]
    fn on_upstream_receive() {
        let config = Config {
            strategy: Strategy::Suffix,
            metadata_key: TOKEN_KEY.into(),
            size: 0,
            remove: false,
        };
        let filter = capture_bytes(config);
        assert_filter_on_downstream_receive_no_change(&filter);
    }

    #[test]
    fn end_capture() {
        let end = Suffix {};
        let mut contents = b"helloabc".to_vec();
        let result = end.capture(&mut contents, 3, false);
        assert_eq!(b"abc".to_vec(), result);
        assert_eq!(b"helloabc".to_vec(), contents);

        let result = end.capture(&mut contents, 3, true);
        assert_eq!(b"abc".to_vec(), result);
        assert_eq!(b"hello".to_vec(), contents);
    }

    #[test]
    fn beginning_capture() {
        let beg = Prefix {};
        let mut contents = b"abchello".to_vec();

        let result = beg.capture(&mut contents, 3, false);
        assert_eq!(b"abc".to_vec(), result);
        assert_eq!(b"abchello".to_vec(), contents);

        let result = beg.capture(&mut contents, 3, true);
        assert_eq!(b"abc".to_vec(), result);
        assert_eq!(b"hello".to_vec(), contents);
    }

    fn assert_end_strategy<F>(filter: &F, key: &str, remove: bool)
    where
        F: Filter + ?Sized,
    {
        let endpoints = vec![EndPoint {
            name: "e1".to_string(),
            address: "127.0.0.1:81".parse().unwrap(),
            connection_ids: vec![],
        }];
        let response = filter
            .on_downstream_receive(DownstreamContext::new(
                endpoints,
                "127.0.0.1:80".parse().unwrap(),
                "helloabc".to_string().into_bytes(),
            ))
            .unwrap();

        if remove {
            assert_eq!(b"hello".to_vec(), response.contents);
        } else {
            assert_eq!(b"helloabc".to_vec(), response.contents);
        }

        let token = response
            .metadata
            .get(key)
            .unwrap()
            .downcast_ref::<Vec<u8>>()
            .unwrap();
        assert_eq!(b"abc", token.as_slice());
    }
}
<commit_msg>Rename CaptureBytes context_key ➡ metadata_key (#133)<commit_after>/*
 * Copyright 2020 Google LLC All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use serde::{Deserialize, Serialize};
use slog::{o, warn, Logger};

use metrics::Metrics;

use crate::extensions::filters::CAPTURED_BYTES;
use crate::extensions::{
    CreateFilterArgs, DownstreamContext, DownstreamResponse, Error, Filter, FilterFactory,
    UpstreamContext, UpstreamResponse,
};

mod metrics;

#[derive(Serialize, Deserialize, Debug)]
/// Strategy to apply for acquiring a set of bytes in the UDP packet
enum Strategy {
    #[serde(rename = "PREFIX")]
    /// Looks for the set of bytes at the beginning of the packet
    Prefix,
    #[serde(rename = "SUFFIX")]
    /// Look for the set of bytes at the end of the packet
    Suffix,
}

#[derive(Serialize, Deserialize, Debug)]
struct Config {
    #[serde(default)]
    strategy: Strategy,
    /// the number of bytes to capture
    #[serde(rename = "size")]
    size: usize,
    /// the key to use when storing the captured bytes in the filter context
    #[serde(rename = "metadataKey")]
    #[serde(default = "default_metadata_key")]
    metadata_key: String,
    /// whether or not to remove the set of the bytes from the packet once captured
    #[serde(default)]
    remove: bool,
}

/// default value for the context key in the Config
fn default_metadata_key() -> String {
    CAPTURED_BYTES.into()
}

impl Default for Strategy {
    fn default() -> Self {
        Strategy::Suffix
    }
}

pub struct CaptureBytesFactory {
    log: Logger,
}

impl CaptureBytesFactory {
    pub fn new(base: &Logger) -> Self {
        CaptureBytesFactory { log: base.clone() }
    }
}

impl FilterFactory for CaptureBytesFactory {
    fn name(&self) -> String {
        "quilkin.extensions.filters.capture_bytes.v1alpha1.CaptureBytes".into()
    }

    fn create_filter(&self, args: CreateFilterArgs) -> Result<Box<dyn Filter>, Error> {
        let config: Config = serde_yaml::to_string(&args.config)
            .and_then(|raw_config| serde_yaml::from_str(raw_config.as_str()))
            .map_err(|err| Error::DeserializeFailed(err.to_string()))?;

        Ok(Box::new(CaptureBytes::new(
            &self.log,
            config,
            Metrics::new(&args.metrics_registry)?,
        )))
    }
}

struct CaptureBytes {
    log: Logger,
    capture: Box<dyn Capture + Sync + Send>,
    /// metrics reporter for this filter.
    metrics: Metrics,
    metadata_key: String,
    size: usize,
    remove: bool,
}

impl CaptureBytes {
    fn new(base: &Logger, config: Config, metrics: Metrics) -> Self {
        let capture: Box<dyn Capture + Sync + Send> = match config.strategy {
            Strategy::Prefix => Box::new(Prefix {}),
            Strategy::Suffix => Box::new(Suffix {}),
        };

        CaptureBytes {
            log: base.new(o!("source" => "extensions::CaptureBytes")),
            capture,
            metrics,
            metadata_key: config.metadata_key,
            size: config.size,
            remove: config.remove,
        }
    }
}

impl Filter for CaptureBytes {
    fn on_downstream_receive(&self, mut ctx: DownstreamContext) -> Option<DownstreamResponse> {
        // if the capture size is bigger than the packet size, then we drop the packet,
        // and occasionally warn
        if ctx.contents.len() < self.size {
            if self.metrics.packets_dropped_total.get() % 1000 == 0 {
                warn!(
                    self.log,
                    "packets are being dropped due to their length being less than {} bytes",
                    self.size; "count" => self.metrics.packets_dropped_total.get()
                );
            }
            self.metrics.packets_dropped_total.inc();
            return None;
        }
        let token = self
            .capture
            .capture(&mut ctx.contents, self.size, self.remove);

        ctx.metadata
            .insert(self.metadata_key.clone(), Box::new(token));

        Some(ctx.into())
    }

    fn on_upstream_receive(&self, ctx: UpstreamContext) -> Option<UpstreamResponse> {
        Some(ctx.into())
    }
}

/// Trait to implement different strategies for capturing packet data
trait Capture {
    /// Capture the packet data from the contents. If remove is true, contents will be altered to
    /// not have the retrieved set of bytes.
    /// Returns the captured bytes.
    fn capture(&self, contents: &mut Vec<u8>, size: usize, remove: bool) -> Vec<u8>;
}

struct Suffix;
impl Capture for Suffix {
    fn capture(&self, contents: &mut Vec<u8>, size: usize, remove: bool) -> Vec<u8> {
        if remove {
            return contents.split_off(contents.len() - size);
        }

        contents
            .iter()
            .skip(contents.len() - size)
            .cloned()
            .collect::<Vec<u8>>()
    }
}

struct Prefix;
impl Capture for Prefix {
    fn capture(&self, contents: &mut Vec<u8>, size: usize, remove: bool) -> Vec<u8> {
        if remove {
            return contents.drain(..size).collect();
        }

        contents.iter().cloned().take(size).collect()
    }
}

#[cfg(test)]
mod tests {
    use prometheus::Registry;
    use serde_yaml::{Mapping, Value};

    use crate::config::{ConnectionConfig, EndPoint};
    use crate::test_utils::{assert_filter_on_downstream_receive_no_change, logger};

    use super::*;

    const TOKEN_KEY: &str = "TOKEN";

    fn capture_bytes(config: Config) -> CaptureBytes {
        CaptureBytes::new(
            &logger(),
            config,
            Metrics::new(&Registry::default()).unwrap(),
        )
    }

    #[test]
    fn factory_valid_config_all() {
        let factory = CaptureBytesFactory::new(&logger());
        let connection = ConnectionConfig::Server { endpoints: vec![] };
        let mut map = Mapping::new();
        map.insert(
            Value::String("strategy".into()),
            Value::String("SUFFIX".into()),
        );
        map.insert(
            Value::String("metadataKey".into()),
            Value::String(TOKEN_KEY.into()),
        );
        map.insert(Value::String("size".into()), Value::Number(3.into()));
        map.insert(Value::String("remove".into()), Value::Bool(true));

        let filter = factory
            .create_filter(CreateFilterArgs::new(
                &connection,
                Some(&Value::Mapping(map)),
            ))
            .unwrap();
        assert_end_strategy(filter.as_ref(), TOKEN_KEY, true);
    }

    #[test]
    fn factory_valid_config_defaults() {
        let factory = CaptureBytesFactory::new(&logger());
        let connection = ConnectionConfig::Server { endpoints: vec![] };
        let mut map = Mapping::new();
        map.insert(Value::String("size".into()), Value::Number(3.into()));
        let filter = factory
            .create_filter(CreateFilterArgs::new(
                &connection,
                Some(&Value::Mapping(map)),
            ))
            .unwrap();
        assert_end_strategy(filter.as_ref(), CAPTURED_BYTES, false);
    }

    #[test]
    fn factory_invalid_config() {
        let factory = CaptureBytesFactory::new(&logger());
        let connection = ConnectionConfig::Server { endpoints: vec![] };
        let mut map = Mapping::new();
        map.insert(Value::String("size".into()), Value::String("WRONG".into()));

        let result = factory.create_filter(CreateFilterArgs::new(
            &connection,
            Some(&Value::Mapping(map)),
        ));
        assert!(result.is_err(), "Should be an error");
    }

    #[test]
    fn on_downstream_receive() {
        let config = Config {
            strategy: Strategy::Suffix,
            metadata_key: TOKEN_KEY.into(),
            size: 3,
            remove: true,
        };
        let filter = capture_bytes(config);
        assert_end_strategy(&filter, TOKEN_KEY, true);
    }

    #[test]
    fn on_downstream_receive_overflow_capture_size() {
        let config = Config {
            strategy: Strategy::Suffix,
            metadata_key: TOKEN_KEY.into(),
            size: 99,
            remove: true,
        };
        let filter = capture_bytes(config);
        let endpoints = vec![EndPoint {
            name: "e1".to_string(),
            address: "127.0.0.1:81".parse().unwrap(),
            connection_ids: vec![],
        }];
        let response = filter.on_downstream_receive(DownstreamContext::new(
            endpoints,
            "127.0.0.1:80".parse().unwrap(),
            "abc".to_string().into_bytes(),
        ));

        assert!(response.is_none());
        let count = filter.metrics.packets_dropped_total.get();
        assert_eq!(1, count);
    }

    #[test]
    fn on_upstream_receive() {
        let config = Config {
            strategy: Strategy::Suffix,
            metadata_key: TOKEN_KEY.into(),
            size: 0,
            remove: false,
        };
        let filter = capture_bytes(config);
        assert_filter_on_downstream_receive_no_change(&filter);
    }

    #[test]
    fn end_capture() {
        let end = Suffix {};
        let mut contents = b"helloabc".to_vec();
        let result = end.capture(&mut contents, 3, false);
        assert_eq!(b"abc".to_vec(), result);
        assert_eq!(b"helloabc".to_vec(), contents);

        let result = end.capture(&mut contents, 3, true);
        assert_eq!(b"abc".to_vec(), result);
        assert_eq!(b"hello".to_vec(), contents);
    }

    #[test]
    fn beginning_capture() {
        let beg = Prefix {};
        let mut contents = b"abchello".to_vec();

        let result = beg.capture(&mut contents, 3, false);
        assert_eq!(b"abc".to_vec(), result);
        assert_eq!(b"abchello".to_vec(), contents);

        let result = beg.capture(&mut contents, 3, true);
        assert_eq!(b"abc".to_vec(), result);
        assert_eq!(b"hello".to_vec(), contents);
    }

    fn assert_end_strategy<F>(filter: &F, key: &str, remove: bool)
    where
        F: Filter + ?Sized,
    {
        let endpoints = vec![EndPoint {
            name: "e1".to_string(),
            address: "127.0.0.1:81".parse().unwrap(),
            connection_ids: vec![],
        }];
        let response = filter
            .on_downstream_receive(DownstreamContext::new(
                endpoints,
                "127.0.0.1:80".parse().unwrap(),
                "helloabc".to_string().into_bytes(),
            ))
            .unwrap();

        if remove {
            assert_eq!(b"hello".to_vec(), response.contents);
        } else {
            assert_eq!(b"helloabc".to_vec(), response.contents);
        }

        let token = response
            .metadata
            .get(key)
            .unwrap()
            .downcast_ref::<Vec<u8>>()
            .unwrap();
        assert_eq!(b"abc", token.as_slice());
    }
}
<|endoftext|>


<commit_before>use super::*;

use ::test;

#[bench]
#[cfg(not(miri))] // Miri does not support benchmarks
fn bench_push_back_100(b: &mut test::Bencher) {
    let mut deq = VecDeque::with_capacity(101);
    b.iter(|| {
        for i in 0..100 {
            deq.push_back(i);
        }
        deq.head = 0;
        deq.tail = 0;
    })
}

#[bench]
#[cfg(not(miri))] // Miri does not support benchmarks
fn bench_push_front_100(b: &mut test::Bencher) {
    let mut deq = VecDeque::with_capacity(101);
    b.iter(|| {
        for i in 0..100 {
            deq.push_front(i);
        }
        deq.head = 0;
        deq.tail = 0;
    })
}

#[bench]
#[cfg(not(miri))] // Miri does not support benchmarks
fn bench_pop_back_100(b: &mut test::Bencher) {
    let mut deq = VecDeque::<i32>::with_capacity(101);

    b.iter(|| {
        deq.head = 100;
        deq.tail = 0;
        while !deq.is_empty() {
            test::black_box(deq.pop_back());
        }
    })
}

#[bench]
#[cfg(not(miri))] // Miri does not support benchmarks
fn bench_pop_front_100(b: &mut test::Bencher) {
    let mut deq = VecDeque::<i32>::with_capacity(101);

    b.iter(|| {
        deq.head = 100;
        deq.tail = 0;
        while !deq.is_empty() {
            test::black_box(deq.pop_front());
        }
    })
}

#[test]
fn test_swap_front_back_remove() {
    fn test(back: bool) {
        // This test checks that every single combination of tail position and length is tested.
        // Capacity 15 should be large enough to cover every case.
        let mut tester = VecDeque::with_capacity(15);
        let usable_cap = tester.capacity();
        let final_len = usable_cap / 2;

        for len in 0..final_len {
            let expected: VecDeque<_> = if back {
                (0..len).collect()
            } else {
                (0..len).rev().collect()
            };
            for tail_pos in 0..usable_cap {
                tester.tail = tail_pos;
                tester.head = tail_pos;
                if back {
                    for i in 0..len * 2 {
                        tester.push_front(i);
                    }
                    for i in 0..len {
                        assert_eq!(tester.swap_remove_back(i), Some(len * 2 - 1 - i));
                    }
                } else {
                    for i in 0..len * 2 {
                        tester.push_back(i);
                    }
                    for i in 0..len {
                        let idx = tester.len() - 1 - i;
                        assert_eq!(tester.swap_remove_front(idx), Some(len * 2 - 1 - i));
                    }
                }
                assert!(tester.tail < tester.cap());
                assert!(tester.head < tester.cap());
                assert_eq!(tester, expected);
            }
        }
    }
    test(true);
    test(false);
}

#[test]
fn test_insert() {
    // This test checks that every single combination of tail position, length, and
    // insertion position is tested. Capacity 15 should be large enough to cover every case.

    let mut tester = VecDeque::with_capacity(15);
    // can't guarantee we got 15, so have to get what we got.
    // 15 would be great, but we will definitely get 2^k - 1, for k >= 4, or else
    // this test isn't covering what it wants to
    let cap = tester.capacity();

    // len is the length *after* insertion
    for len in 1..cap {
        // 0, 1, 2, .., len - 1
        let expected = (0..).take(len).collect::<VecDeque<_>>();
        for tail_pos in 0..cap {
            for to_insert in 0..len {
                tester.tail = tail_pos;
                tester.head = tail_pos;
                for i in 0..len {
                    if i != to_insert {
                        tester.push_back(i);
                    }
                }
                tester.insert(to_insert, to_insert);
                assert!(tester.tail < tester.cap());
                assert!(tester.head < tester.cap());
                assert_eq!(tester, expected);
            }
        }
    }
}

#[test]
fn test_remove() {
    // This test checks that every single combination of tail position, length, and
    // removal position is tested. Capacity 15 should be large enough to cover every case.

    let mut tester = VecDeque::with_capacity(15);
    // can't guarantee we got 15, so have to get what we got.
    // 15 would be great, but we will definitely get 2^k - 1, for k >= 4, or else
    // this test isn't covering what it wants to
    let cap = tester.capacity();

    // len is the length *after* removal
    for len in 0..cap - 1 {
        // 0, 1, 2, .., len - 1
        let expected = (0..).take(len).collect::<VecDeque<_>>();
        for tail_pos in 0..cap {
            for to_remove in 0..=len {
                tester.tail = tail_pos;
                tester.head = tail_pos;
                for i in 0..len {
                    if i == to_remove {
                        tester.push_back(1234);
                    }
                    tester.push_back(i);
                }
                if to_remove == len {
                    tester.push_back(1234);
                }
                tester.remove(to_remove);
                assert!(tester.tail < tester.cap());
                assert!(tester.head < tester.cap());
                assert_eq!(tester, expected);
            }
        }
    }
}

#[test]
fn test_drain() {
    let mut tester: VecDeque<usize> = VecDeque::with_capacity(7);

    let cap = tester.capacity();
    for len in 0..=cap {
        for tail in 0..=cap {
            for drain_start in 0..=len {
                for drain_end in drain_start..=len {
                    tester.tail = tail;
                    tester.head = tail;
                    for i in 0..len {
                        tester.push_back(i);
                    }

                    // Check that we drain the correct values
                    let drained: VecDeque<_> = tester.drain(drain_start..drain_end).collect();
                    let drained_expected: VecDeque<_> = (drain_start..drain_end).collect();
                    assert_eq!(drained, drained_expected);

                    // We shouldn't have changed the capacity or made the
                    // head or tail out of bounds
                    assert_eq!(tester.capacity(), cap);
                    assert!(tester.tail < tester.cap());
                    assert!(tester.head < tester.cap());

                    // We should see the correct values in the VecDeque
                    let expected: VecDeque<_> = (0..drain_start)
                        .chain(drain_end..len)
                        .collect();
                    assert_eq!(expected, tester);
                }
            }
        }
    }
}

#[test]
fn test_shrink_to_fit() {
    // This test checks that every single combination of head and tail position,
    // is tested. Capacity 15 should be large enough to cover every case.

    let mut tester = VecDeque::with_capacity(15);
    // can't guarantee we got 15, so have to get what we got.
    // 15 would be great, but we will definitely get 2^k - 1, for k >= 4, or else
    // this test isn't covering what it wants to
    let cap = tester.capacity();
    tester.reserve(63);
    let max_cap = tester.capacity();

    for len in 0..=cap {
        // 0, 1, 2, .., len - 1
        let expected = (0..).take(len).collect::<VecDeque<_>>();
        for tail_pos in 0..=max_cap {
            tester.tail = tail_pos;
            tester.head = tail_pos;
            tester.reserve(63);
            for i in 0..len {
                tester.push_back(i);
            }
            tester.shrink_to_fit();
            assert!(tester.capacity() <= cap);
            assert!(tester.tail < tester.cap());
            assert!(tester.head < tester.cap());
            assert_eq!(tester, expected);
        }
    }
}

#[test]
fn test_split_off() {
    // This test checks that every single combination of tail position, length, and
    // split position is tested. Capacity 15 should be large enough to cover every case.

    let mut tester = VecDeque::with_capacity(15);
    // can't guarantee we got 15, so have to get what we got.
    // 15 would be great, but we will definitely get 2^k - 1, for k >= 4, or else
    // this test isn't covering what it wants to
    let cap = tester.capacity();

    // len is the length *before* splitting
    for len in 0..cap {
        // index to split at
        for at in 0..=len {
            // 0, 1, 2, .., at - 1 (may be empty)
            let expected_self = (0..).take(at).collect::<VecDeque<_>>();
            // at, at + 1, .., len - 1 (may be empty)
            let expected_other = (at..).take(len - at).collect::<VecDeque<_>>();

            for tail_pos in 0..cap {
                tester.tail = tail_pos;
                tester.head = tail_pos;
                for i in 0..len {
                    tester.push_back(i);
                }
                let result = tester.split_off(at);
                assert!(tester.tail < tester.cap());
                assert!(tester.head < tester.cap());
                assert!(result.tail < result.cap());
                assert!(result.head < result.cap());
                assert_eq!(tester, expected_self);
                assert_eq!(result, expected_other);
            }
        }
    }
}

#[test]
fn test_from_vec() {
    use crate::vec::Vec;
    for cap in 0..35 {
        for len in 0..=cap {
            let mut vec = Vec::with_capacity(cap);
            vec.extend(0..len);

            let vd = VecDeque::from(vec.clone());
            assert!(vd.cap().is_power_of_two());
            assert_eq!(vd.len(), vec.len());
            assert!(vd.into_iter().eq(vec));
        }
    }
}

#[test]
fn test_vec_from_vecdeque() {
    use crate::vec::Vec;

    fn create_vec_and_test_convert(capacity: usize, offset: usize, len: usize) {
        let mut vd = VecDeque::with_capacity(capacity);
        for _ in 0..offset {
            vd.push_back(0);
            vd.pop_front();
        }
        vd.extend(0..len);

        let vec: Vec<_> = Vec::from(vd.clone());
        assert_eq!(vec.len(), vd.len());
        assert!(vec.into_iter().eq(vd));
    }

    #[cfg(not(miri))] // Miri is too slow
    let max_pwr = 7;
    #[cfg(miri)]
    let max_pwr = 5;

    for cap_pwr in 0..max_pwr {
        // Make capacity as a (2^x)-1, so that the ring size is 2^x
        let cap = (2i32.pow(cap_pwr) - 1) as usize;

        // In these cases there is enough free space to solve it with copies
        for len in 0..((cap + 1) / 2) {
            // Test contiguous cases
            for offset in 0..(cap - len) {
                create_vec_and_test_convert(cap, offset, len)
            }

            // Test cases where block at end of buffer is bigger than block at start
            for offset in (cap - len)..(cap - (len / 2)) {
                create_vec_and_test_convert(cap, offset, len)
            }

            // Test cases where block at start of buffer is bigger than block at end
            for offset in (cap - (len / 2))..cap {
                create_vec_and_test_convert(cap, offset, len)
            }
        }

        // Now there's not (necessarily) space to straighten the ring with simple copies,
        // the ring will use swapping when:
        // (cap + 1 - offset) > (cap + 1 - len) && (len - (cap + 1 - offset)) > (cap + 1 - len))
        //  right block size  >   free space    &&      left block size       >    free space
        for len in ((cap + 1) / 2)..cap {
            // Test contiguous cases
            for offset in 0..(cap - len) {
                create_vec_and_test_convert(cap, offset, len)
            }

            // Test cases where block at end of buffer is bigger than block at start
            for offset in (cap - len)..(cap - (len / 2)) {
                create_vec_and_test_convert(cap, offset, len)
            }

            // Test cases where block at start of buffer is bigger than block at end
            for offset in (cap - (len / 2))..cap {
                create_vec_and_test_convert(cap, offset, len)
            }
        }
    }
}

#[test]
fn issue_53529() {
    use crate::boxed::Box;

    let mut dst = VecDeque::new();
    dst.push_front(Box::new(1));
    dst.push_front(Box::new(2));
    assert_eq!(*dst.pop_back().unwrap(), 1);

    let mut src = VecDeque::new();
    src.push_front(Box::new(2));
    dst.append(&mut src);
    for a in dst {
        assert_eq!(*a, 2);
    }
}
<commit_msg>Add tests for VecDeque clone_from<commit_after>use super::*;

use ::test;

#[bench]
#[cfg(not(miri))] // Miri does not support benchmarks
fn bench_push_back_100(b: &mut test::Bencher) {
    let mut deq = VecDeque::with_capacity(101);
    b.iter(|| {
        for i in 0..100 {
            deq.push_back(i);
        }
        deq.head = 0;
        deq.tail = 0;
    })
}

#[bench]
#[cfg(not(miri))] // Miri does not support benchmarks
fn bench_push_front_100(b: &mut test::Bencher) {
    let mut deq = VecDeque::with_capacity(101);
    b.iter(|| {
        for i in 0..100 {
            deq.push_front(i);
        }
        deq.head = 0;
        deq.tail = 0;
    })
}

#[bench]
#[cfg(not(miri))] // Miri does not support benchmarks
fn bench_pop_back_100(b: &mut test::Bencher) {
    let mut deq = VecDeque::<i32>::with_capacity(101);

    b.iter(|| {
        deq.head = 100;
        deq.tail = 0;
        while !deq.is_empty() {
            test::black_box(deq.pop_back());
        }
    })
}

#[bench]
#[cfg(not(miri))] // Miri does not support benchmarks
fn bench_pop_front_100(b: &mut test::Bencher) {
    let mut deq = VecDeque::<i32>::with_capacity(101);

    b.iter(|| {
        deq.head = 100;
        deq.tail = 0;
        while !deq.is_empty() {
            test::black_box(deq.pop_front());
        }
    })
}

#[test]
fn test_swap_front_back_remove() {
    fn test(back: bool) {
        // This test checks that every single combination of tail position and length is tested.
        // Capacity 15 should be large enough to cover every case.
        let mut tester = VecDeque::with_capacity(15);
        let usable_cap = tester.capacity();
        let final_len = usable_cap / 2;

        for len in 0..final_len {
            let expected: VecDeque<_> = if back {
                (0..len).collect()
            } else {
                (0..len).rev().collect()
            };
            for tail_pos in 0..usable_cap {
                tester.tail = tail_pos;
                tester.head = tail_pos;
                if back {
                    for i in 0..len * 2 {
                        tester.push_front(i);
                    }
                    for i in 0..len {
                        assert_eq!(tester.swap_remove_back(i), Some(len * 2 - 1 - i));
                    }
                } else {
                    for i in 0..len * 2 {
                        tester.push_back(i);
                    }
                    for i in 0..len {
                        let idx = tester.len() - 1 - i;
                        assert_eq!(tester.swap_remove_front(idx), Some(len * 2 - 1 - i));
                    }
                }
                assert!(tester.tail < tester.cap());
                assert!(tester.head < tester.cap());
                assert_eq!(tester, expected);
            }
        }
    }
    test(true);
    test(false);
}

#[test]
fn test_insert() {
    // This test checks that every single combination of tail position, length, and
    // insertion position is tested. Capacity 15 should be large enough to cover every case.

    let mut tester = VecDeque::with_capacity(15);
    // can't guarantee we got 15, so have to get what we got.
    // 15 would be great, but we will definitely get 2^k - 1, for k >= 4, or else
    // this test isn't covering what it wants to
    let cap = tester.capacity();

    // len is the length *after* insertion
    for len in 1..cap {
        // 0, 1, 2, .., len - 1
        let expected = (0..).take(len).collect::<VecDeque<_>>();
        for tail_pos in 0..cap {
            for to_insert in 0..len {
                tester.tail = tail_pos;
                tester.head = tail_pos;
                for i in 0..len {
                    if i != to_insert {
                        tester.push_back(i);
                    }
                }
                tester.insert(to_insert, to_insert);
                assert!(tester.tail < tester.cap());
                assert!(tester.head < tester.cap());
                assert_eq!(tester, expected);
            }
        }
    }
}

#[test]
fn test_remove() {
    // This test checks that every single combination of tail position, length, and
    // removal position is tested. Capacity 15 should be large enough to cover every case.

    let mut tester = VecDeque::with_capacity(15);
    // can't guarantee we got 15, so have to get what we got.
    // 15 would be great, but we will definitely get 2^k - 1, for k >= 4, or else
    // this test isn't covering what it wants to
    let cap = tester.capacity();

    // len is the length *after* removal
    for len in 0..cap - 1 {
        // 0, 1, 2, .., len - 1
        let expected = (0..).take(len).collect::<VecDeque<_>>();
        for tail_pos in 0..cap {
            for to_remove in 0..=len {
                tester.tail = tail_pos;
                tester.head = tail_pos;
                for i in 0..len {
                    if i == to_remove {
                        tester.push_back(1234);
                    }
                    tester.push_back(i);
                }
                if to_remove == len {
                    tester.push_back(1234);
                }
                tester.remove(to_remove);
                assert!(tester.tail < tester.cap());
                assert!(tester.head < tester.cap());
                assert_eq!(tester, expected);
            }
        }
    }
}

#[test]
fn test_drain() {
    let mut tester: VecDeque<usize> = VecDeque::with_capacity(7);

    let cap = tester.capacity();
    for len in 0..=cap {
        for tail in 0..=cap {
            for drain_start in 0..=len {
                for drain_end in drain_start..=len {
                    tester.tail = tail;
                    tester.head = tail;
                    for i in 0..len {
                        tester.push_back(i);
                    }

                    // Check that we drain the correct values
                    let drained: VecDeque<_> = tester.drain(drain_start..drain_end).collect();
                    let drained_expected: VecDeque<_> = (drain_start..drain_end).collect();
                    assert_eq!(drained, drained_expected);

                    // We shouldn't have changed the capacity or made the
                    // head or tail out of bounds
                    assert_eq!(tester.capacity(), cap);
                    assert!(tester.tail < tester.cap());
                    assert!(tester.head < tester.cap());

                    // We should see the correct values in the VecDeque
                    let expected: VecDeque<_> = (0..drain_start)
                        .chain(drain_end..len)
                        .collect();
                    assert_eq!(expected, tester);
                }
            }
        }
    }
}

#[test]
fn test_shrink_to_fit() {
    // This test checks that every single combination of head and tail position,
    // is tested. Capacity 15 should be large enough to cover every case.

    let mut tester = VecDeque::with_capacity(15);
    // can't guarantee we got 15, so have to get what we got.
    // 15 would be great, but we will definitely get 2^k - 1, for k >= 4, or else
    // this test isn't covering what it wants to
    let cap = tester.capacity();
    tester.reserve(63);
    let max_cap = tester.capacity();

    for len in 0..=cap {
        // 0, 1, 2, .., len - 1
        let expected = (0..).take(len).collect::<VecDeque<_>>();
        for tail_pos in 0..=max_cap {
            tester.tail = tail_pos;
            tester.head = tail_pos;
            tester.reserve(63);
            for i in 0..len {
                tester.push_back(i);
            }
            tester.shrink_to_fit();
            assert!(tester.capacity() <= cap);
            assert!(tester.tail < tester.cap());
            assert!(tester.head < tester.cap());
            assert_eq!(tester, expected);
        }
    }
}

#[test]
fn test_split_off() {
    // This test checks that every single combination of tail position, length, and
    // split position is tested. Capacity 15 should be large enough to cover every case.

    let mut tester = VecDeque::with_capacity(15);
    // can't guarantee we got 15, so have to get what we got.
    // 15 would be great, but we will definitely get 2^k - 1, for k >= 4, or else
    // this test isn't covering what it wants to
    let cap = tester.capacity();

    // len is the length *before* splitting
    for len in 0..cap {
        // index to split at
        for at in 0..=len {
            // 0, 1, 2, .., at - 1 (may be empty)
            let expected_self = (0..).take(at).collect::<VecDeque<_>>();
            // at, at + 1, .., len - 1 (may be empty)
            let expected_other = (at..).take(len - at).collect::<VecDeque<_>>();

            for tail_pos in 0..cap {
                tester.tail = tail_pos;
                tester.head = tail_pos;
                for i in 0..len {
                    tester.push_back(i);
                }
                let result = tester.split_off(at);
                assert!(tester.tail < tester.cap());
                assert!(tester.head < tester.cap());
                assert!(result.tail < result.cap());
                assert!(result.head < result.cap());
                assert_eq!(tester, expected_self);
                assert_eq!(result, expected_other);
            }
        }
    }
}

#[test]
fn test_from_vec() {
    use crate::vec::Vec;
    for cap in 0..35 {
        for len in 0..=cap {
            let mut vec = Vec::with_capacity(cap);
            vec.extend(0..len);

            let vd = VecDeque::from(vec.clone());
            assert!(vd.cap().is_power_of_two());
            assert_eq!(vd.len(), vec.len());
            assert!(vd.into_iter().eq(vec));
        }
    }
}

#[test]
fn test_vec_from_vecdeque() {
    use crate::vec::Vec;

    fn create_vec_and_test_convert(capacity: usize, offset: usize, len: usize) {
        let mut vd = VecDeque::with_capacity(capacity);
        for _ in 0..offset {
            vd.push_back(0);
            vd.pop_front();
        }
        vd.extend(0..len);

        let vec: Vec<_> = Vec::from(vd.clone());
        assert_eq!(vec.len(), vd.len());
        assert!(vec.into_iter().eq(vd));
    }

    #[cfg(not(miri))] // Miri is too slow
    let max_pwr = 7;
    #[cfg(miri)]
    let max_pwr = 5;

    for cap_pwr in 0..max_pwr {
        // Make capacity as a (2^x)-1, so that the ring size is 2^x
        let cap = (2i32.pow(cap_pwr) - 1) as usize;

        // In these cases there is enough free space to solve it with copies
        for len in 0..((cap + 1) / 2) {
            // Test contiguous cases
            for offset in 0..(cap - len) {
                create_vec_and_test_convert(cap, offset, len)
            }

            // Test cases where block at end of buffer is bigger than block at start
            for offset in (cap - len)..(cap - (len / 2)) {
                create_vec_and_test_convert(cap, offset, len)
            }

            // Test cases where block at start of buffer is bigger than block at end
            for offset in (cap - (len / 2))..cap {
                create_vec_and_test_convert(cap, offset, len)
            }
        }

        // Now there's not (necessarily) space to straighten the ring with simple copies,
        // the ring will use swapping when:
        // (cap + 1 - offset) > (cap + 1 - len) && (len - (cap + 1 - offset)) > (cap + 1 - len))
        //  right block size  >   free space    &&      left block size       >    free space
        for len in ((cap + 1) / 2)..cap {
            // Test contiguous cases
            for offset in 0..(cap - len) {
                create_vec_and_test_convert(cap, offset, len)
            }

            // Test cases where block at end of buffer is bigger than block at start
            for offset in (cap - len)..(cap - (len / 2)) {
                create_vec_and_test_convert(cap, offset, len)
            }

            // Test cases where block at start of buffer is bigger than block at end
            for offset in (cap - (len / 2))..cap {
                create_vec_and_test_convert(cap, offset, len)
            }
        }
    }
}

#[test]
fn test_clone_from() {
    let m = vec![1; 8];
    let n = vec![2; 12];
    for pfv in 0..8 {
        for pfu in 0..8 {
            for longer in 0..2 {
                let (vr, ur) = if longer == 0 { (&m, &n) } else { (&n, &m) };
                let mut v = VecDeque::from(vr.clone());
                for _ in 0..pfv {
                    v.push_front(1);
                }
                let mut u = VecDeque::from(ur.clone());
                for _ in 0..pfu {
                    u.push_front(2);
                }
                v.clone_from(&u);
                assert_eq!(&v, &u);
            }
        }
    }
}

#[test]
fn issue_53529() {
    use crate::boxed::Box;

    let mut dst = VecDeque::new();
    dst.push_front(Box::new(1));
    dst.push_front(Box::new(2));
    assert_eq!(*dst.pop_back().unwrap(), 1);

    let mut src = VecDeque::new();
    src.push_front(Box::new(2));
    dst.append(&mut src);
    for a in dst {
        assert_eq!(*a, 2);
    }
}
<|endoftext|>


<commit_before>use crate::{CompileError, Config, Syntax};

use std::path::{Path, PathBuf};
use std::str::FromStr;

use quote::ToTokens;

pub struct TemplateInput<'a> {
    pub ast: &'a syn::DeriveInput,
    pub config: &'a Config<'a>,
    pub syntax: &'a Syntax<'a>,
    pub source: Source,
    pub print: Print,
    pub escaper: &'a str,
    pub ext: Option<String>,
    pub parent: Option<&'a syn::Type>,
    pub path: PathBuf,
}

impl<'a> TemplateInput<'a> {
    /// Extract the template metadata from the `DeriveInput` structure. This
    /// mostly recovers the data for the `TemplateInput` fields from the
    /// `template()` attribute list fields; it also finds the of the `_parent`
    /// field, if any.
    pub fn new<'n>(
        ast: &'n syn::DeriveInput,
        config: &'n Config,
    ) -> Result<TemplateInput<'n>, CompileError> {
        // Check that an attribute called `template()` exists and that it is
        // the proper type (list).
        let meta = ast
            .attrs
            .iter()
            .find_map(|attr| match attr.parse_meta() {
                Ok(m) => {
                    if m.path().is_ident("template") {
                        Some(Ok(m))
                    } else {
                        None
                    }
                }
                Err(e) => Some(Err(format!("unable to parse attribute: {}", e).into())),
            })
            .unwrap_or(Err(CompileError::Static("no attribute 'template' found")))?;

        let meta_list = match meta {
            syn::Meta::List(inner) => inner,
            _ => return Err("attribute 'template' has incorrect type".into()),
        };

        // Loop over the meta attributes and find everything that we
        // understand. Raise panics if something is not right.
        // `source` contains an enum that can represent `path` or `source`.
        let mut source = None;
        let mut print = Print::None;
        let mut escaping = None;
        let mut ext = None;
        let mut syntax = None;
        for item in meta_list.nested {
            let pair = match item {
                syn::NestedMeta::Meta(syn::Meta::NameValue(ref pair)) => pair,
                _ => {
                    return Err(format!(
                        "unsupported attribute argument {:?}",
                        item.to_token_stream()
                    )
                    .into())
                }
            };

            if pair.path.is_ident("path") {
                if let syn::Lit::Str(ref s) = pair.lit {
                    if source.is_some() {
                        return Err("must specify 'source' or 'path', not both".into());
                    }
                    source = Some(Source::Path(s.value()));
                } else {
                    return Err("template path must be string literal".into());
                }
            } else if pair.path.is_ident("source") {
                if let syn::Lit::Str(ref s) = pair.lit {
                    if source.is_some() {
                        return Err("must specify 'source' or 'path', not both".into());
                    }
                    source = Some(Source::Source(s.value()));
                } else {
                    return Err("template source must be string literal".into());
                }
            } else if pair.path.is_ident("print") {
                if let syn::Lit::Str(ref s) = pair.lit {
                    print = s.value().parse()?;
                } else {
                    return Err("print value must be string literal".into());
                }
            } else if pair.path.is_ident("escape") {
                if let syn::Lit::Str(ref s) = pair.lit {
                    escaping = Some(s.value());
                } else {
                    return Err("escape value must be string literal".into());
                }
            } else if pair.path.is_ident("ext") {
                if let syn::Lit::Str(ref s) = pair.lit {
                    ext = Some(s.value());
                } else {
                    return Err("ext value must be string literal".into());
                }
            } else if pair.path.is_ident("syntax") {
                if let syn::Lit::Str(ref s) = pair.lit {
                    syntax = Some(s.value())
                } else {
                    return Err("syntax value must be string literal".into());
                }
            } else {
                return Err(format!(
                    "unsupported attribute key '{}' found",
                    pair.path.to_token_stream()
                )
                .into());
            }
        }

        // Validate the `source` and `ext` value together, since they are
        // related. In case `source` was used instead of `path`, the value
        // of `ext` is merged into a synthetic `path` value here.
        let source = source.expect("template path or source not found in attributes");
        let path = match (&source, &ext) {
            (&Source::Path(ref path), None) => config.find_template(path, None)?,
            (&Source::Source(_), Some(ext)) => PathBuf::from(format!("{}.{}", ast.ident, ext)),
            (&Source::Path(_), Some(_)) => {
                return Err("'ext' attribute cannot be used with 'path' attribute".into())
            }
            (&Source::Source(_), None) => {
                return Err("must include 'ext' attribute when using 'source' attribute".into())
            }
        };

        // Check to see if a `_parent` field was defined on the context
        // struct, and store the type for it for use in the code generator.
        let parent = match ast.data {
            syn::Data::Struct(syn::DataStruct {
                fields: syn::Fields::Named(ref fields),
                ..
            }) => fields
                .named
                .iter()
                .find(|f| f.ident.as_ref().filter(|name| *name == "_parent").is_some())
                .map(|f| &f.ty),
            _ => None,
        };

        if parent.is_some() {
            eprint!(
                "   --> in struct {}\n   = use of deprecated field '_parent'\n",
                ast.ident
            );
        }

        // Validate syntax
        let syntax = syntax.map_or_else(
            || Ok(config.syntaxes.get(config.default_syntax).unwrap()),
            |s| {
                config.syntaxes.get(&s).ok_or_else(|| {
                    CompileError::String(format!("attribute syntax {} not exist", s))
                })
            },
        )?;

        // Match extension against defined output formats

        let extension = escaping.unwrap_or_else(|| {
            path.extension()
                .map(|s| s.to_str().unwrap())
                .unwrap_or("")
                .to_string()
        });

        let mut escaper = None;
        for (extensions, path) in &config.escapers {
            if extensions.contains(&extension) {
                escaper = Some(path);
                break;
            }
        }

        let escaper = escaper.ok_or_else(|| {
            CompileError::String(format!("no escaper defined for extension '{}'", extension,))
        })?;

        Ok(TemplateInput {
            ast,
            config,
            syntax,
            source,
            print,
            escaper,
            ext,
            parent,
            path,
        })
    }

    pub fn extension(&self) -> Option<&str> {
        extension(&self.path)
    }
}

fn extension(path: &Path) -> Option<&str> {
    let ext = path.extension().map(|s| s.to_str().unwrap())?;

    const JINJA_EXTENSIONS: [&str; 3] = ["j2", "jinja", "jinja2"];
    if JINJA_EXTENSIONS.contains(&ext) {
        Path::new(path.file_stem().unwrap())
            .extension()
            .map(|s| s.to_str().unwrap())
            .or(Some(ext))
    } else {
        Some(ext)
    }
}

pub enum Source {
    Path(String),
    Source(String),
}

#[derive(PartialEq)]
pub enum Print {
    All,
    Ast,
    Code,
    None,
}

impl FromStr for Print {
    type Err = CompileError;

    fn from_str(s: &str) -> Result<Print, Self::Err> {
        use self::Print::*;
        Ok(match s {
            "all" => All,
            "ast" => Ast,
            "code" => Code,
            "none" => None,
            v => return Err(format!("invalid value for print option: {}", v,).into()),
        })
    }
}
<commit_msg>Added extension tests<commit_after>use crate::{CompileError, Config, Syntax};

use std::path::{Path, PathBuf};
use std::str::FromStr;

use quote::ToTokens;

pub struct TemplateInput<'a> {
    pub ast: &'a syn::DeriveInput,
    pub config: &'a Config<'a>,
    pub syntax: &'a Syntax<'a>,
    pub source: Source,
    pub print: Print,
    pub escaper: &'a str,
    pub ext: Option<String>,
    pub parent: Option<&'a syn::Type>,
    pub path: PathBuf,
}

impl<'a> TemplateInput<'a> {
    /// Extract the template metadata from the `DeriveInput` structure. This
    /// mostly recovers the data for the `TemplateInput` fields from the
    /// `template()` attribute list fields; it also finds the of the `_parent`
    /// field, if any.
    pub fn new<'n>(
        ast: &'n syn::DeriveInput,
        config: &'n Config,
    ) -> Result<TemplateInput<'n>, CompileError> {
        // Check that an attribute called `template()` exists and that it is
        // the proper type (list).
        let meta = ast
            .attrs
            .iter()
            .find_map(|attr| match attr.parse_meta() {
                Ok(m) => {
                    if m.path().is_ident("template") {
                        Some(Ok(m))
                    } else {
                        None
                    }
                }
                Err(e) => Some(Err(format!("unable to parse attribute: {}", e).into())),
            })
            .unwrap_or(Err(CompileError::Static("no attribute 'template' found")))?;

        let meta_list = match meta {
            syn::Meta::List(inner) => inner,
            _ => return Err("attribute 'template' has incorrect type".into()),
        };

        // Loop over the meta attributes and find everything that we
        // understand. Raise panics if something is not right.
        // `source` contains an enum that can represent `path` or `source`.
        let mut source = None;
        let mut print = Print::None;
        let mut escaping = None;
        let mut ext = None;
        let mut syntax = None;
        for item in meta_list.nested {
            let pair = match item {
                syn::NestedMeta::Meta(syn::Meta::NameValue(ref pair)) => pair,
                _ => {
                    return Err(format!(
                        "unsupported attribute argument {:?}",
                        item.to_token_stream()
                    )
                    .into())
                }
            };

            if pair.path.is_ident("path") {
                if let syn::Lit::Str(ref s) = pair.lit {
                    if source.is_some() {
                        return Err("must specify 'source' or 'path', not both".into());
                    }
                    source = Some(Source::Path(s.value()));
                } else {
                    return Err("template path must be string literal".into());
                }
            } else if pair.path.is_ident("source") {
                if let syn::Lit::Str(ref s) = pair.lit {
                    if source.is_some() {
                        return Err("must specify 'source' or 'path', not both".into());
                    }
                    source = Some(Source::Source(s.value()));
                } else {
                    return Err("template source must be string literal".into());
                }
            } else if pair.path.is_ident("print") {
                if let syn::Lit::Str(ref s) = pair.lit {
                    print = s.value().parse()?;
                } else {
                    return Err("print value must be string literal".into());
                }
            } else if pair.path.is_ident("escape") {
                if let syn::Lit::Str(ref s) = pair.lit {
                    escaping = Some(s.value());
                } else {
                    return Err("escape value must be string literal".into());
                }
            } else if pair.path.is_ident("ext") {
                if let syn::Lit::Str(ref s) = pair.lit {
                    ext = Some(s.value());
                } else {
                    return Err("ext value must be string literal".into());
                }
            } else if pair.path.is_ident("syntax") {
                if let syn::Lit::Str(ref s) = pair.lit {
                    syntax = Some(s.value())
                } else {
                    return Err("syntax value must be string literal".into());
                }
            } else {
                return Err(format!(
                    "unsupported attribute key '{}' found",
                    pair.path.to_token_stream()
                )
                .into());
            }
        }

        // Validate the `source` and `ext` value together, since they are
        // related. In case `source` was used instead of `path`, the value
        // of `ext` is merged into a synthetic `path` value here.
        let source = source.expect("template path or source not found in attributes");
        let path = match (&source, &ext) {
            (&Source::Path(ref path), None) => config.find_template(path, None)?,
            (&Source::Source(_), Some(ext)) => PathBuf::from(format!("{}.{}", ast.ident, ext)),
            (&Source::Path(_), Some(_)) => {
                return Err("'ext' attribute cannot be used with 'path' attribute".into())
            }
            (&Source::Source(_), None) => {
                return Err("must include 'ext' attribute when using 'source' attribute".into())
            }
        };

        // Check to see if a `_parent` field was defined on the context
        // struct, and store the type for it for use in the code generator.
        let parent = match ast.data {
            syn::Data::Struct(syn::DataStruct {
                fields: syn::Fields::Named(ref fields),
                ..
            }) => fields
                .named
                .iter()
                .find(|f| f.ident.as_ref().filter(|name| *name == "_parent").is_some())
                .map(|f| &f.ty),
            _ => None,
        };

        if parent.is_some() {
            eprint!(
                "   --> in struct {}\n   = use of deprecated field '_parent'\n",
                ast.ident
            );
        }

        // Validate syntax
        let syntax = syntax.map_or_else(
            || Ok(config.syntaxes.get(config.default_syntax).unwrap()),
            |s| {
                config.syntaxes.get(&s).ok_or_else(|| {
                    CompileError::String(format!("attribute syntax {} not exist", s))
                })
            },
        )?;

        // Match extension against defined output formats

        let extension = escaping.unwrap_or_else(|| {
            path.extension()
                .map(|s| s.to_str().unwrap())
                .unwrap_or("")
                .to_string()
        });

        let mut escaper = None;
        for (extensions, path) in &config.escapers {
            if extensions.contains(&extension) {
                escaper = Some(path);
                break;
            }
        }

        let escaper = escaper.ok_or_else(|| {
            CompileError::String(format!("no escaper defined for extension '{}'", extension,))
        })?;

        Ok(TemplateInput {
            ast,
            config,
            syntax,
            source,
            print,
            escaper,
            ext,
            parent,
            path,
        })
    }

    pub fn extension(&self) -> Option<&str> {
        extension(&self.path)
    }
}

fn extension(path: &Path) -> Option<&str> {
    let ext = path.extension().map(|s| s.to_str().unwrap())?;

    const JINJA_EXTENSIONS: [&str; 3] = ["j2", "jinja", "jinja2"];
    if JINJA_EXTENSIONS.contains(&ext) {
        Path::new(path.file_stem().unwrap())
            .extension()
            .map(|s| s.to_str().unwrap())
            .or(Some(ext))
    } else {
        Some(ext)
    }
}

pub enum Source {
    Path(String),
    Source(String),
}

#[derive(PartialEq)]
pub enum Print {
    All,
    Ast,
    Code,
    None,
}

impl FromStr for Print {
    type Err = CompileError;

    fn from_str(s: &str) -> Result<Print, Self::Err> {
        use self::Print::*;
        Ok(match s {
            "all" => All,
            "ast" => Ast,
            "code" => Code,
            "none" => None,
            v => return Err(format!("invalid value for print option: {}", v,).into()),
        })
    }
}

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

    #[test]
    fn test_ext() {
        assert_eq!(extension(Path::new("foo-bar.txt")), Some("txt"));
        assert_eq!(extension(Path::new("foo-bar.html")), Some("html"));
        assert_eq!(extension(Path::new("foo-bar.unknown")), Some("unknown"));

        assert_eq!(extension(Path::new("foo/bar/baz.txt")), Some("txt"));
        assert_eq!(extension(Path::new("foo/bar/baz.html")), Some("html"));
        assert_eq!(extension(Path::new("foo/bar/baz.unknown")), Some("unknown"));
    }

    #[test]
    fn test_double_ext() {
        assert_eq!(extension(Path::new("foo-bar.html.txt")), Some("txt"));
        assert_eq!(extension(Path::new("foo-bar.txt.html")), Some("html"));
        assert_eq!(extension(Path::new("foo-bar.txt.unknown")), Some("unknown"));

        assert_eq!(extension(Path::new("foo/bar/baz.html.txt")), Some("txt"));
        assert_eq!(extension(Path::new("foo/bar/baz.txt.html")), Some("html"));
        assert_eq!(
            extension(Path::new("foo/bar/baz.txt.unknown")),
            Some("unknown")
        );
    }

    #[test]
    fn test_skip_jinja_ext() {
        assert_eq!(extension(Path::new("foo-bar.html.j2")), Some("html"));
        assert_eq!(extension(Path::new("foo-bar.html.jinja")), Some("html"));
        assert_eq!(extension(Path::new("foo-bar.html.jinja2")), Some("html"));

        assert_eq!(extension(Path::new("foo/bar/baz.txt.j2")), Some("txt"));
        assert_eq!(extension(Path::new("foo/bar/baz.txt.jinja")), Some("txt"));
        assert_eq!(extension(Path::new("foo/bar/baz.txt.jinja2")), Some("txt"));
    }

    #[test]
    fn test_only_jinja_ext() {
        assert_eq!(extension(Path::new("foo-bar.j2")), Some("j2"));
        assert_eq!(extension(Path::new("foo-bar.jinja")), Some("jinja"));
        assert_eq!(extension(Path::new("foo-bar.jinja2")), Some("jinja2"));
    }
}
<|endoftext|>