// Copyright 2022 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! The loopback device.

use alloc::vec::Vec;
use core::convert::Infallible as Never;
use core::fmt::Debug;
use derivative::Derivative;

use lock_order::lock::{OrderedLockAccess, OrderedLockRef};
use log::trace;
use net_types::ethernet::Mac;
use net_types::ip::{Ipv4, Ipv6, Mtu};
use netstack3_base::sync::Mutex;
use netstack3_base::{
    AnyDevice, BroadcastIpExt, CoreTimerContext, Device, DeviceIdAnyCompatContext, DeviceIdContext,
    FrameDestination, RecvFrameContext, RecvIpFrameMeta, ResourceCounterContext, SendFrameError,
    SendFrameErrorReason, SendableFrameMeta, StrongDeviceIdentifier, TimerContext,
    WeakDeviceIdentifier,
};
use netstack3_ip::IpPacketDestination;
use packet::{Buf, Buffer as _, BufferMut, Serializer};
use packet_formats::ethernet::{
    EtherType, EthernetFrame, EthernetFrameBuilder, EthernetFrameLengthCheck, EthernetIpExt,
};

use crate::internal::base::{
    DeviceCounters, DeviceLayerTypes, DeviceReceiveFrameSpec, EthernetDeviceCounters,
};
use crate::internal::id::{BaseDeviceId, BasePrimaryDeviceId, BaseWeakDeviceId, WeakDeviceId};
use crate::internal::queue::rx::{
    ReceiveDequeFrameContext, ReceiveQueue, ReceiveQueueState, ReceiveQueueTypes,
};
use crate::internal::queue::tx::{
    BufVecU8Allocator, TransmitQueue, TransmitQueueHandler, TransmitQueueState,
};
use crate::internal::queue::{DequeueState, TransmitQueueFrameError};
use crate::internal::socket::{
    DeviceSocketHandler, DeviceSocketMetadata, DeviceSocketSendTypes, EthernetHeaderParams,
    ReceivedFrame,
};
use crate::internal::state::{DeviceStateSpec, IpLinkDeviceState};

/// The MAC address corresponding to the loopback interface.
const LOOPBACK_MAC: Mac = Mac::UNSPECIFIED;

/// A weak device ID identifying a loopback device.
///
/// This device ID is like [`WeakDeviceId`] but specifically for loopback
/// devices.
///
/// [`WeakDeviceId`]: crate::device::WeakDeviceId
pub type LoopbackWeakDeviceId<BT> = BaseWeakDeviceId<LoopbackDevice, BT>;

/// A strong device ID identifying a loopback device.
///
/// This device ID is like [`DeviceId`] but specifically for loopback devices.
///
/// [`DeviceId`]: crate::device::DeviceId
pub type LoopbackDeviceId<BT> = BaseDeviceId<LoopbackDevice, BT>;

/// The primary reference for a loopback device.
pub type LoopbackPrimaryDeviceId<BT> = BasePrimaryDeviceId<LoopbackDevice, BT>;

/// Loopback device domain.
#[derive(Copy, Clone)]
pub enum LoopbackDevice {}

impl Device for LoopbackDevice {}

impl DeviceStateSpec for LoopbackDevice {
    type Link<BT: DeviceLayerTypes> = LoopbackDeviceState<WeakDeviceId<BT>>;
    type External<BT: DeviceLayerTypes> = BT::LoopbackDeviceState;
    type CreationProperties = LoopbackCreationProperties;
    type Counters = EthernetDeviceCounters;
    type TimerId<D: WeakDeviceIdentifier> = Never;

    fn new_link_state<
        CC: CoreTimerContext<Self::TimerId<CC::WeakDeviceId>, BC> + DeviceIdContext<Self>,
        BC: DeviceLayerTypes + TimerContext,
    >(
        _bindings_ctx: &mut BC,
        _self_id: CC::WeakDeviceId,
        LoopbackCreationProperties { mtu }: Self::CreationProperties,
    ) -> Self::Link<BC> {
        LoopbackDeviceState {
            counters: Default::default(),
            mtu,
            rx_queue: Default::default(),
            tx_queue: Default::default(),
        }
    }

    const IS_LOOPBACK: bool = true;
    const DEBUG_TYPE: &'static str = "Loopback";
}

/// Properties used to create a loopback device.
#[derive(Debug)]
pub struct LoopbackCreationProperties {
    /// The device's MTU.
    pub mtu: Mtu,
}

/// State for a loopback device.
pub struct LoopbackDeviceState<D: WeakDeviceIdentifier> {
    /// Loopback device counters.
    pub counters: EthernetDeviceCounters,
    /// The MTU this device was created with (immutable).
    pub mtu: Mtu,
    /// Loopback device receive queue.
    pub rx_queue: ReceiveQueue<LoopbackRxQueueMeta<D>, Buf<Vec<u8>>>,
    /// Loopback device transmit queue.
    pub tx_queue: TransmitQueue<LoopbackTxQueueMeta<D>, Buf<Vec<u8>>, BufVecU8Allocator>,
}

#[derive(Derivative)]
#[derivative(Default(bound = ""))]
/// Metadata associated with a frame in the Loopback TX queue.
pub struct LoopbackTxQueueMeta<D: WeakDeviceIdentifier> {
    /// Device that should be used to deliver the packet. If not set then the
    /// packet delivered as if it came from the loopback device.
    target_device: Option<D>,
}

/// Metadata associated with a frame in the Loopback RX queue.
pub struct LoopbackRxQueueMeta<D: WeakDeviceIdentifier> {
    /// Device that should be used to deliver the packet. If not set then the
    /// packet delivered as if it came from the loopback device.
    target_device: Option<D>,
}

impl<D: WeakDeviceIdentifier> From<LoopbackTxQueueMeta<D>> for LoopbackRxQueueMeta<D> {
    fn from(LoopbackTxQueueMeta { target_device }: LoopbackTxQueueMeta<D>) -> Self {
        Self { target_device }
    }
}

impl<BT: DeviceLayerTypes>
    OrderedLockAccess<ReceiveQueueState<LoopbackRxQueueMeta<WeakDeviceId<BT>>, Buf<Vec<u8>>>>
    for IpLinkDeviceState<LoopbackDevice, BT>
{
    type Lock = Mutex<ReceiveQueueState<LoopbackRxQueueMeta<WeakDeviceId<BT>>, Buf<Vec<u8>>>>;
    fn ordered_lock_access(&self) -> OrderedLockRef<'_, Self::Lock> {
        OrderedLockRef::new(&self.link.rx_queue.queue)
    }
}

impl<BT: DeviceLayerTypes>
    OrderedLockAccess<DequeueState<LoopbackRxQueueMeta<WeakDeviceId<BT>>, Buf<Vec<u8>>>>
    for IpLinkDeviceState<LoopbackDevice, BT>
{
    type Lock = Mutex<DequeueState<LoopbackRxQueueMeta<WeakDeviceId<BT>>, Buf<Vec<u8>>>>;
    fn ordered_lock_access(&self) -> OrderedLockRef<'_, Self::Lock> {
        OrderedLockRef::new(&self.link.rx_queue.deque)
    }
}

impl<BT: DeviceLayerTypes>
    OrderedLockAccess<
        TransmitQueueState<LoopbackTxQueueMeta<WeakDeviceId<BT>>, Buf<Vec<u8>>, BufVecU8Allocator>,
    > for IpLinkDeviceState<LoopbackDevice, BT>
{
    type Lock = Mutex<
        TransmitQueueState<LoopbackTxQueueMeta<WeakDeviceId<BT>>, Buf<Vec<u8>>, BufVecU8Allocator>,
    >;
    fn ordered_lock_access(&self) -> OrderedLockRef<'_, Self::Lock> {
        OrderedLockRef::new(&self.link.tx_queue.queue)
    }
}

impl<BT: DeviceLayerTypes>
    OrderedLockAccess<DequeueState<LoopbackTxQueueMeta<WeakDeviceId<BT>>, Buf<Vec<u8>>>>
    for IpLinkDeviceState<LoopbackDevice, BT>
{
    type Lock = Mutex<DequeueState<LoopbackTxQueueMeta<WeakDeviceId<BT>>, Buf<Vec<u8>>>>;
    fn ordered_lock_access(&self) -> OrderedLockRef<'_, Self::Lock> {
        OrderedLockRef::new(&self.link.tx_queue.deque)
    }
}

impl DeviceSocketSendTypes for LoopbackDevice {
    /// When `None`, data will be sent as a raw Ethernet frame without any
    /// system-applied headers.
    type Metadata = Option<EthernetHeaderParams>;
}

impl<CC, BC> ReceiveDequeFrameContext<LoopbackDevice, BC> for CC
where
    CC: DeviceIdContext<LoopbackDevice>
        + ResourceCounterContext<Self::DeviceId, EthernetDeviceCounters>
        + ReceiveQueueTypes<
            LoopbackDevice,
            BC,
            Meta = LoopbackRxQueueMeta<<CC as DeviceIdContext<AnyDevice>>::WeakDeviceId>,
        >,
    // Loopback needs to deliver messages to `AnyDevice`.
    CC: DeviceIdAnyCompatContext<LoopbackDevice>
        + RecvFrameContext<RecvIpFrameMeta<<CC as DeviceIdContext<AnyDevice>>::DeviceId, Ipv4>, BC>
        + RecvFrameContext<RecvIpFrameMeta<<CC as DeviceIdContext<AnyDevice>>::DeviceId, Ipv6>, BC>
        + ResourceCounterContext<<CC as DeviceIdContext<AnyDevice>>::DeviceId, DeviceCounters>
        + DeviceSocketHandler<AnyDevice, BC>,
    CC::Buffer: BufferMut + Debug,
    BC: DeviceLayerTypes,
{
    fn handle_frame(
        &mut self,
        bindings_ctx: &mut BC,
        device_id: &Self::DeviceId,
        rx_meta: Self::Meta,
        mut buf: Self::Buffer,
    ) {
        let (frame, whole_body) =
            match buf.parse_with_view::<_, EthernetFrame<_>>(EthernetFrameLengthCheck::NoCheck) {
                Err(e) => {
                    self.increment(&device_id.clone().into(), |counters: &DeviceCounters| {
                        &counters.recv_parse_error
                    });
                    trace!("dropping invalid ethernet frame over loopback: {:?}", e);
                    return;
                }
                Ok(e) => e,
            };

        let target_device: <CC as DeviceIdContext<AnyDevice>>::DeviceId =
            match rx_meta.target_device.map(|d| d.upgrade()) {
                // This is a packet that should be delivered on `target_device`.
                Some(Some(dev)) => dev,

                // `target_device` is gone. Drop the packet.
                Some(None) => return,

                // This is a packet sent to the loopback device.
                None => device_id.clone().into(),
            };

        self.increment(&target_device, |counters: &DeviceCounters| &counters.recv_frame);

        let frame_dest = FrameDestination::from_dest(frame.dst_mac(), Mac::UNSPECIFIED);
        let ethertype = frame.ethertype();

        DeviceSocketHandler::<AnyDevice, _>::handle_frame(
            self,
            bindings_ctx,
            &target_device,
            ReceivedFrame::from_ethernet(frame, frame_dest).into(),
            whole_body,
        );

        match ethertype {
            Some(EtherType::Ipv4) => {
                self.increment(&target_device, |counters: &DeviceCounters| {
                    &counters.recv_ipv4_delivered
                });
                self.receive_frame(
                    bindings_ctx,
                    RecvIpFrameMeta::<_, Ipv4>::new(target_device, Some(frame_dest)),
                    buf,
                );
            }
            Some(EtherType::Ipv6) => {
                self.increment(&target_device, |counters: &DeviceCounters| {
                    &counters.recv_ipv6_delivered
                });
                self.receive_frame(
                    bindings_ctx,
                    RecvIpFrameMeta::<_, Ipv6>::new(target_device, Some(frame_dest)),
                    buf,
                );
            }
            Some(ethertype @ (EtherType::Arp | EtherType::Other(_))) => {
                self.increment(device_id, |counters: &EthernetDeviceCounters| {
                    &counters.recv_unsupported_ethertype
                });
                trace!("not handling loopback frame of type {:?}", ethertype)
            }
            None => {
                self.increment(device_id, |counters: &EthernetDeviceCounters| {
                    &counters.recv_no_ethertype
                });
                trace!("dropping ethernet frame without ethertype");
            }
        }
    }
}

impl<CC, BC> SendableFrameMeta<CC, BC>
    for DeviceSocketMetadata<LoopbackDevice, <CC as DeviceIdContext<LoopbackDevice>>::DeviceId>
where
    CC: TransmitQueueHandler<
            LoopbackDevice,
            BC,
            Meta = LoopbackTxQueueMeta<<CC as DeviceIdContext<AnyDevice>>::WeakDeviceId>,
        > + ResourceCounterContext<<CC as DeviceIdContext<LoopbackDevice>>::DeviceId, DeviceCounters>
        + DeviceIdContext<AnyDevice>,
    BC: DeviceLayerTypes,
{
    fn send_meta<S>(
        self,
        core_ctx: &mut CC,
        bindings_ctx: &mut BC,
        body: S,
    ) -> Result<(), SendFrameError<S>>
    where
        S: Serializer,
        S::Buffer: BufferMut,
    {
        let Self { device_id, metadata } = self;
        let tx_meta = LoopbackTxQueueMeta::default();
        match metadata {
            Some(EthernetHeaderParams { dest_addr, protocol }) => send_as_ethernet_frame_to_dst(
                core_ctx,
                bindings_ctx,
                &device_id,
                body,
                protocol,
                dest_addr,
                LoopbackTxQueueMeta::default(),
            ),
            None => send_ethernet_frame(core_ctx, bindings_ctx, &device_id, body, tx_meta),
        }
    }
}

/// Sends an IP frame `packet` over `device_id`.
pub fn send_ip_frame<CC, BC, I, S>(
    core_ctx: &mut CC,
    bindings_ctx: &mut BC,
    device_id: &<CC as DeviceIdContext<LoopbackDevice>>::DeviceId,
    destination: IpPacketDestination<I, &<CC as DeviceIdContext<AnyDevice>>::DeviceId>,
    packet: S,
) -> Result<(), SendFrameError<S>>
where
    CC: TransmitQueueHandler<
            LoopbackDevice,
            BC,
            Meta = LoopbackTxQueueMeta<<CC as DeviceIdContext<AnyDevice>>::WeakDeviceId>,
        > + ResourceCounterContext<<CC as DeviceIdContext<LoopbackDevice>>::DeviceId, DeviceCounters>
        + DeviceIdContext<AnyDevice>,
    BC: DeviceLayerTypes,
    I: EthernetIpExt + BroadcastIpExt,
    S: Serializer,
    S::Buffer: BufferMut,
{
    core_ctx.increment(device_id, DeviceCounters::send_frame::<I>);

    let target_device = match destination {
        IpPacketDestination::Loopback(device) => Some(device.downgrade()),
        IpPacketDestination::Broadcast(_)
        | IpPacketDestination::Multicast(_)
        | IpPacketDestination::Neighbor(_) => None,
    };
    send_as_ethernet_frame_to_dst(
        core_ctx,
        bindings_ctx,
        device_id,
        packet,
        I::ETHER_TYPE,
        LOOPBACK_MAC,
        LoopbackTxQueueMeta { target_device },
    )
}

fn send_as_ethernet_frame_to_dst<CC, BC, S>(
    core_ctx: &mut CC,
    bindings_ctx: &mut BC,
    device_id: &<CC as DeviceIdContext<LoopbackDevice>>::DeviceId,
    packet: S,
    protocol: EtherType,
    dst_mac: Mac,
    meta: LoopbackTxQueueMeta<<CC as DeviceIdContext<AnyDevice>>::WeakDeviceId>,
) -> Result<(), SendFrameError<S>>
where
    CC: TransmitQueueHandler<
            LoopbackDevice,
            BC,
            Meta = LoopbackTxQueueMeta<<CC as DeviceIdContext<AnyDevice>>::WeakDeviceId>,
        > + ResourceCounterContext<<CC as DeviceIdContext<LoopbackDevice>>::DeviceId, DeviceCounters>
        + DeviceIdContext<AnyDevice>,
    BC: DeviceLayerTypes,
    S: Serializer,
    S::Buffer: BufferMut,
{
    /// The minimum length of bodies of Ethernet frames sent over the loopback
    /// device.
    ///
    /// Use zero since the frames are never sent out a physical device, so it
    /// doesn't matter if they are shorter than would be required.
    const MIN_BODY_LEN: usize = 0;

    let frame = packet.encapsulate(EthernetFrameBuilder::new(
        LOOPBACK_MAC,
        dst_mac,
        protocol,
        MIN_BODY_LEN,
    ));

    send_ethernet_frame(core_ctx, bindings_ctx, device_id, frame, meta)
        .map_err(|err| err.into_inner())
}

fn send_ethernet_frame<CC, BC, S>(
    core_ctx: &mut CC,
    bindings_ctx: &mut BC,
    device_id: &<CC as DeviceIdContext<LoopbackDevice>>::DeviceId,
    frame: S,
    meta: LoopbackTxQueueMeta<<CC as DeviceIdContext<AnyDevice>>::WeakDeviceId>,
) -> Result<(), SendFrameError<S>>
where
    CC: TransmitQueueHandler<
            LoopbackDevice,
            BC,
            Meta = LoopbackTxQueueMeta<<CC as DeviceIdContext<AnyDevice>>::WeakDeviceId>,
        > + ResourceCounterContext<<CC as DeviceIdContext<LoopbackDevice>>::DeviceId, DeviceCounters>
        + DeviceIdContext<AnyDevice>,
    S: Serializer,
    S::Buffer: BufferMut,
    BC: DeviceLayerTypes,
{
    core_ctx.increment(device_id, |counters: &DeviceCounters| &counters.send_total_frames);
    match TransmitQueueHandler::<LoopbackDevice, _>::queue_tx_frame(
        core_ctx,
        bindings_ctx,
        device_id,
        meta,
        frame,
    ) {
        Ok(()) => {
            core_ctx.increment(device_id, |counters: &DeviceCounters| &counters.send_frame);
            Ok(())
        }
        Err(TransmitQueueFrameError::NoQueue(err)) => {
            unreachable!("loopback never fails to send a frame: {err:?}")
        }
        Err(TransmitQueueFrameError::QueueFull(serializer)) => {
            core_ctx.increment(device_id, |counters: &DeviceCounters| &counters.send_queue_full);
            Err(SendFrameError { serializer, error: SendFrameErrorReason::QueueFull })
        }
        Err(TransmitQueueFrameError::SerializeError(err)) => {
            core_ctx
                .increment(device_id, |counters: &DeviceCounters| &counters.send_serialize_error);
            Err(err.err_into())
        }
    }
}

impl DeviceReceiveFrameSpec for LoopbackDevice {
    // Loopback never receives frames from bindings, so make it impossible to
    // instantiate it.
    type FrameMetadata<D> = Never;
}