EasyTier/easytier/src/instance/runtime_host/tun_common.rs

use std::{
    any::Any,
    sync::{Arc, OnceLock},
};

use easytier_core::{host::packet::HostPacketReceiver, instance::CorePacketPlane};
use tokio::{sync::Mutex, task::JoinSet};

use super::MagicDnsRuntime;
use crate::{
    common::{error::Error, global_ctx::ArcGlobalCtx},
    instance::{shared_virtual_nic::ArcSharedVirtualNicRegistry, virtual_nic::NicCtx},
};

pub(super) async fn create_nic_ctx(
    global_ctx: ArcGlobalCtx,
    packet_plane: Arc<CorePacketPlane>,
    receiver: Arc<Mutex<HostPacketReceiver>>,
    close_notifier: Arc<tokio::sync::Notify>,
    registry: ArcSharedVirtualNicRegistry,
) -> Result<NicCtx, Error> {
    #[cfg(not(mobile))]
    if global_ctx.get_flags().dev_name.is_empty() {
        return Ok(NicCtx::new(
            global_ctx,
            packet_plane,
            receiver,
            close_notifier,
        ));
    }

    let member_id = global_ctx.get_id();
    NicCtx::new_shared(
        global_ctx,
        packet_plane,
        receiver,
        close_notifier,
        registry,
        member_id,
    )
    .await
}

struct NicCtxContainer {
    _nic_ctx: Option<Box<dyn Any + Send>>,
    magic_dns: MagicDnsRuntime,
}

impl NicCtxContainer {
    fn new(nic_ctx: NicCtx, magic_dns: MagicDnsRuntime) -> Self {
        Self {
            _nic_ctx: Some(Box::new(nic_ctx)),
            magic_dns,
        }
    }

    fn packet_drain(tasks: JoinSet<()>) -> Self {
        Self {
            _nic_ctx: Some(Box::new(tasks)),
            magic_dns: MagicDnsRuntime::default(),
        }
    }
}

#[derive(Clone)]
pub(super) struct TunNicState {
    nic_ctx: Arc<Mutex<Option<NicCtxContainer>>>,
    receiver: Arc<OnceLock<Arc<Mutex<HostPacketReceiver>>>>,
}

impl TunNicState {
    pub(super) fn empty() -> Self {
        Self {
            nic_ctx: Arc::new(Mutex::new(None)),
            receiver: Arc::new(OnceLock::new()),
        }
    }

    pub(super) fn install_receiver(&self, receiver: HostPacketReceiver) -> anyhow::Result<()> {
        self.receiver
            .set(Arc::new(Mutex::new(receiver)))
            .map_err(|_| anyhow::anyhow!("native packet receiver is already installed"))
    }

    pub(super) fn receiver(&self) -> Arc<Mutex<HostPacketReceiver>> {
        self.receiver
            .get()
            .expect("packet receiver must be installed before preparing TUN")
            .clone()
    }

    pub(super) async fn stop(&self) {
        let mut old = self.nic_ctx.lock().await.take();
        if let Some(nic) = old.as_mut() {
            nic.magic_dns.stop().await;
        }
        drop(old);
    }

    pub(super) async fn drain(&self) {
        self.stop().await;
        let receiver = self.receiver();
        let mut tasks = JoinSet::new();
        tasks.spawn(async move {
            let mut receiver = receiver.lock().await;
            while let Some(packet) = receiver.recv().await {
                tracing::trace!(?packet, "discarded packet without a native interface");
            }
        });
        self.nic_ctx
            .lock()
            .await
            .replace(NicCtxContainer::packet_drain(tasks));
    }

    pub(super) async fn install(&self, nic: NicCtx, magic_dns: MagicDnsRuntime) {
        self.stop().await;
        self.nic_ctx
            .lock()
            .await
            .replace(NicCtxContainer::new(nic, magic_dns));
    }
}