v2

vnt-core/src/utils/device_id.rs

@@ -0,0 +1,29 @@
+use anyhow::Context;
+use std::fs;
+use std::path::Path;
+pub fn get_device_id() -> anyhow::Result<String> {
+    match machine_uid::get() {
+        Ok(id) => return Ok(id),
+        Err(e) => {
+            log::warn!("Failed to get system ID: {}. Using fallback.", e);
+        }
+    }
+
+    get_fallback_id()
+}
+fn get_fallback_id() -> anyhow::Result<String> {
+    let path = Path::new("device_id");
+
+    if let Ok(content) = fs::read_to_string(path) {
+        let id = content.trim();
+        if !id.is_empty() {
+            return Ok(id.to_string());
+        }
+    }
+
+    let new_id = uuid::Uuid::new_v4().to_string();
+
+    fs::write(path, &new_id).context("Failed to write device_id file")?;
+
+    Ok(new_id)
+}

vnt-core/src/utils/dns_query.rs

@@ -0,0 +1,251 @@
+use anyhow::{Context, anyhow};
+use dns_parser::{Builder, Packet, QueryClass, QueryType, RData, ResponseCode};
+use rand::seq::SliceRandom;
+use std::io;
+use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr};
+use std::str::FromStr;
+use std::time::Duration;
+use tokio::net::UdpSocket;
+
+use rust_p2p_core::socket::LocalInterface;
+
+pub async fn dns_query_txt(
+    domain: &str,
+    mut name_servers: Vec<String>,
+    default_interface: &Option<LocalInterface>,
+) -> io::Result<Vec<String>> {
+    let mut err: Option<io::Error> = None;
+    if name_servers.is_empty() {
+        name_servers.push("223.5.5.5:53".into());
+        name_servers.push("114.114.114.114:53".into());
+    }
+    for name_server in name_servers {
+        match txt_dns(domain, name_server, default_interface).await {
+            Ok(addr) => {
+                if !addr.is_empty() {
+                    return Ok(addr);
+                }
+            }
+            Err(e) => {
+                err.replace(e);
+            }
+        }
+        continue;
+    }
+    if let Some(e) = err {
+        Err(e)
+    } else {
+        Err(io::Error::other(format!("DNS query failed {domain:?}")))
+    }
+}
+pub async fn dns_query_one(
+    domain: &str,
+    name_servers: &Vec<String>,
+    default_interface: &Option<LocalInterface>,
+) -> anyhow::Result<SocketAddr> {
+    let mut vec = dns_query_all(domain, name_servers, default_interface).await?;
+    vec.shuffle(&mut rand::rng());
+    vec.pop().context("DNS query failed")
+}
+pub async fn dns_query_all(
+    domain: &str,
+    name_servers: &Vec<String>,
+    default_interface: &Option<LocalInterface>,
+) -> anyhow::Result<Vec<SocketAddr>> {
+    match SocketAddr::from_str(domain) {
+        Ok(addr) => Ok(vec![addr]),
+        Err(_) => {
+            if name_servers.is_empty() {
+                let addrs: Vec<SocketAddr> = tokio::net::lookup_host(domain)
+                    .await
+                    .map_err(|e| io::Error::other(format!("DNS query failed: {domain:?},{e:?}")))?
+                    .collect();
+                return Ok(addrs);
+            }
+
+            let mut err: Option<io::Error> = None;
+            for name_server in name_servers {
+                let end_index = domain
+                    .rfind(':')
+                    .ok_or_else(|| io::Error::other(format!("not port: {domain:?}")))?;
+                let host = &domain[..end_index];
+                let port = u16::from_str(&domain[end_index + 1..])
+                    .map_err(|_| io::Error::other(format!("not port: {domain:?}")))?;
+                let th1 = {
+                    let host = host.to_string();
+                    let name_server = name_server.clone();
+                    let default_interface = default_interface.clone();
+                    tokio::spawn(a_dns(host, name_server, default_interface.clone()))
+                };
+                let th2 = {
+                    let host = host.to_string();
+                    let name_server = name_server.clone();
+                    let default_interface = default_interface.clone();
+                    tokio::spawn(aaaa_dns(host, name_server, default_interface.clone()))
+                };
+                let mut addr = Vec::new();
+                match th1.await? {
+                    Ok(rs) => {
+                        for ip in rs {
+                            addr.push(SocketAddr::new(ip.into(), port));
+                        }
+                    }
+                    Err(e) => {
+                        err.replace(e);
+                    }
+                }
+                match th2.await? {
+                    Ok(rs) => {
+                        for ip in rs {
+                            addr.push(SocketAddr::new(ip.into(), port));
+                        }
+                    }
+                    Err(e) => {
+                        if addr.is_empty() {
+                            err.replace(e);
+                            continue;
+                        }
+                    }
+                }
+                if addr.is_empty() {
+                    continue;
+                }
+                return Ok(addr);
+            }
+            if let Some(e) = err {
+                Err(e.into())
+            } else {
+                Err(anyhow!("DNS query failed {domain:?}"))
+            }
+        }
+    }
+}
+
+async fn query<'a>(
+    udp: &UdpSocket,
+    domain: &str,
+    name_server: SocketAddr,
+    record_type: QueryType,
+    buf: &'a mut [u8],
+) -> io::Result<Packet<'a>> {
+    let mut builder = Builder::new_query(1, true);
+    builder.add_question(domain, false, record_type, QueryClass::IN);
+    let packet = builder.build().unwrap();
+
+    udp.connect(name_server).await?;
+    let mut count = 0;
+    let len = loop {
+        udp.send(&packet).await?;
+
+        match tokio::time::timeout(Duration::from_secs(3), udp.recv(buf)).await {
+            Ok(len) => {
+                break len?;
+            }
+            Err(_) => {
+                count += 1;
+                if count < 3 {
+                    continue;
+                }
+                Err(io::Error::other(format!("DNS {name_server:?} recv error ")))?
+            }
+        };
+    };
+
+    let pkt = Packet::parse(&buf[..len]).map_err(|e| {
+        io::Error::other(format!(
+            "domain {domain:?} DNS {name_server:?} data error: {e}"
+        ))
+    })?;
+    if pkt.header.response_code != ResponseCode::NoError {
+        return Err(io::Error::other(format!(
+            "response_code {} DNS {:?} domain {:?}",
+            pkt.header.response_code, name_server, domain
+        )));
+    }
+    if pkt.answers.is_empty() {
+        return Err(io::Error::other(format!(
+            "No records received DNS {name_server:?} domain {domain:?}"
+        )));
+    }
+
+    Ok(pkt)
+}
+
+pub async fn txt_dns(
+    domain: &str,
+    name_server: String,
+    default_interface: &Option<LocalInterface>,
+) -> io::Result<Vec<String>> {
+    let name_server: SocketAddr = name_server
+        .parse()
+        .map_err(|e| io::Error::other(format!("dns {name_server} is error :{e:?}")))?;
+    let udp = bind_udp(name_server, default_interface)?;
+    let mut buf = vec![0u8; 65536];
+    let message = query(&udp, domain, name_server, QueryType::TXT, &mut buf).await?;
+    let mut rs = Vec::new();
+    for record in message.answers {
+        if let RData::TXT(txt) = record.data {
+            for x in txt.iter() {
+                let txt = std::str::from_utf8(x)
+                    .map_err(|_| io::Error::other("record type txt is not string"))?;
+                rs.push(txt.to_string());
+            }
+        }
+    }
+    Ok(rs)
+}
+
+fn bind_udp(
+    name_server: SocketAddr,
+    default_interface: &Option<LocalInterface>,
+) -> io::Result<UdpSocket> {
+    let addr: SocketAddr = if name_server.is_ipv4() {
+        "0.0.0.0:0"
+            .parse()
+            .expect("valid IPv4 socket address literal")
+    } else {
+        "[::]:0".parse().expect("valid IPv6 socket address literal")
+    };
+    let socket = rust_p2p_core::socket::bind_udp(addr, default_interface.as_ref())?;
+    UdpSocket::from_std(socket.into())
+}
+
+pub async fn a_dns(
+    domain: String,
+    name_server: String,
+    default_interface: Option<LocalInterface>,
+) -> io::Result<Vec<Ipv4Addr>> {
+    let name_server: SocketAddr = name_server
+        .parse()
+        .map_err(|e| io::Error::other(format!("dns {name_server} is error :{e:?}")))?;
+    let udp = bind_udp(name_server, &default_interface)?;
+    let mut buf = vec![0u8; 65536];
+    let message = query(&udp, &domain, name_server, QueryType::A, &mut buf).await?;
+    let mut rs = Vec::new();
+    for record in message.answers {
+        if let RData::A(a) = record.data {
+            rs.push(a.0);
+        }
+    }
+    Ok(rs)
+}
+
+pub async fn aaaa_dns(
+    domain: String,
+    name_server: String,
+    default_interface: Option<LocalInterface>,
+) -> io::Result<Vec<Ipv6Addr>> {
+    let name_server: SocketAddr = name_server
+        .parse()
+        .map_err(|e| io::Error::other(format!("dns {name_server} is error :{e:?}")))?;
+    let udp = bind_udp(name_server, &default_interface)?;
+    let mut buf = vec![0u8; 65536];
+    let message = query(&udp, &domain, name_server, QueryType::AAAA, &mut buf).await?;
+    let mut rs = Vec::new();
+    for record in message.answers {
+        if let RData::AAAA(a) = record.data {
+            rs.push(a.0);
+        }
+    }
+    Ok(rs)
+}

vnt-core/src/utils/mod.rs

@@ -0,0 +1,11 @@
+pub mod device_id;
+pub(crate) mod dns_query;
+pub mod task_control;
+pub(crate) mod time {
+    pub fn now_ts_ms() -> i64 {
+        std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .unwrap_or_default()
+            .as_millis() as i64
+    }
+}

vnt-core/src/utils/task_control.rs

@@ -0,0 +1,255 @@
+use parking_lot::Mutex;
+use std::collections::HashMap;
+use std::future::Future;
+use std::sync::{Arc, Weak};
+use tokio::sync::Notify;
+use tokio::task::{Id, JoinHandle};
+
+struct TaskGroupState {
+    stopped: bool,
+    tasks: HashMap<Id, JoinHandle<()>>,
+}
+
+struct TaskGroupInner {
+    state: Mutex<TaskGroupState>,
+    all_stopped_notify: Notify,
+}
+
+impl TaskGroupInner {
+    fn new() -> Self {
+        Self {
+            state: Mutex::new(TaskGroupState {
+                stopped: false,
+                tasks: HashMap::new(),
+            }),
+            all_stopped_notify: Notify::new(),
+        }
+    }
+
+    fn spawn<F>(self: &Arc<Self>, f: F) -> Option<Id>
+    where
+        F: Future + Send + 'static,
+        F::Output: Send + 'static,
+    {
+        let mut state = self.state.lock();
+        if state.stopped {
+            return None;
+        }
+
+        let guard = TaskGuard {
+            inner: Arc::downgrade(self),
+        };
+
+        let handle = tokio::spawn(async move {
+            let _guard = guard;
+            f.await;
+        });
+
+        let task_id = handle.id();
+        state.tasks.insert(task_id, handle);
+        Some(task_id)
+    }
+
+    fn stop(&self) {
+        let mut state = self.state.lock();
+        state.stopped = true;
+        for (_, handle) in state.tasks.drain() {
+            handle.abort();
+        }
+    }
+
+    fn is_stopped(&self) -> bool {
+        self.state.lock().stopped
+    }
+
+    fn remove_task(&self, task_id: Id) {
+        let all_stopped = {
+            let mut state = self.state.lock();
+            state.tasks.remove(&task_id);
+            if state.tasks.is_empty() {
+                state.stopped = true;
+                true
+            } else {
+                false
+            }
+        };
+        if all_stopped {
+            self.all_stopped_notify.notify_waiters();
+        }
+    }
+
+    async fn abort_task(&self, task_id: Id) {
+        let handle = self.state.lock().tasks.remove(&task_id);
+        if let Some(handle) = handle {
+            handle.abort();
+            _ = handle.await;
+        }
+    }
+
+    async fn join_all(&self) {
+        let tasks = std::mem::take(&mut self.state.lock().tasks);
+        for (_, h) in tasks {
+            let _ = h.await;
+        }
+    }
+
+    fn all_tasks_stopped(&self) -> bool {
+        let state = self.state.lock();
+        state.stopped && state.tasks.is_empty()
+    }
+}
+
+impl Drop for TaskGroupInner {
+    fn drop(&mut self) {
+        self.stop();
+    }
+}
+
+struct TaskGuard {
+    inner: Weak<TaskGroupInner>,
+}
+
+impl Drop for TaskGuard {
+    fn drop(&mut self) {
+        if let Some(inner) = self.inner.upgrade() {
+            let task_id = tokio::task::id();
+            inner.remove_task(task_id);
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct TaskGroup {
+    inner: Arc<TaskGroupInner>,
+}
+
+impl TaskGroup {
+    fn new() -> Self {
+        Self {
+            inner: Arc::new(TaskGroupInner::new()),
+        }
+    }
+
+    pub fn stop(&self) {
+        self.inner.stop();
+    }
+
+    pub fn is_stopped(&self) -> bool {
+        self.inner.is_stopped()
+    }
+
+    pub fn spawn<F>(&self, f: F) -> SubTask
+    where
+        F: Future + Send + 'static,
+        F::Output: Send + 'static,
+    {
+        match self.inner.spawn(f) {
+            Some(task_id) => SubTask::new(task_id, Arc::downgrade(&self.inner)),
+            None => SubTask::empty(),
+        }
+    }
+
+    pub async fn join_all(&self) {
+        self.inner.join_all().await;
+    }
+
+    pub async fn wait_all_stopped(&self) {
+        loop {
+            if self.inner.all_tasks_stopped() {
+                return;
+            }
+            self.inner.all_stopped_notify.notified().await;
+        }
+    }
+}
+
+pub struct SubTask {
+    task_id: Option<Id>,
+    inner: Weak<TaskGroupInner>,
+}
+
+impl SubTask {
+    fn new(task_id: Id, inner: Weak<TaskGroupInner>) -> Self {
+        Self {
+            task_id: Some(task_id),
+            inner,
+        }
+    }
+
+    fn empty() -> Self {
+        Self {
+            task_id: None,
+            inner: Weak::new(),
+        }
+    }
+
+    pub async fn stop(&self) {
+        if let Some(task_id) = self.task_id
+            && let Some(inner) = self.inner.upgrade()
+        {
+            inner.abort_task(task_id).await;
+        }
+    }
+
+    pub fn is_running(&self) -> bool {
+        if let Some(task_id) = self.task_id
+            && let Some(inner) = self.inner.upgrade()
+        {
+            return inner.state.lock().tasks.contains_key(&task_id);
+        }
+        false
+    }
+
+    pub fn id(&self) -> Option<Id> {
+        self.task_id
+    }
+}
+#[derive(Clone, Default)]
+pub struct TaskGroupManager {
+    task_group: Arc<Mutex<Option<TaskGroup>>>,
+}
+
+impl TaskGroupManager {
+    pub fn new() -> Self {
+        TaskGroupManager::default()
+    }
+
+    pub fn is_running(&self) -> bool {
+        self.task_group.lock().is_some()
+    }
+
+    pub fn is_stopped(&self) -> bool {
+        self.task_group.lock().is_none()
+    }
+
+    pub fn create_task(&self) -> anyhow::Result<(TaskGroup, TaskGroupGuard)> {
+        let mut guard = self.task_group.lock();
+        if guard.is_some() {
+            anyhow::bail!("运行中")
+        }
+
+        let task_group = TaskGroup::new();
+        guard.replace(task_group.clone());
+        let stop_guard = TaskGroupGuard {
+            task_group: self.task_group.clone(),
+        };
+        Ok((task_group, stop_guard))
+    }
+
+    pub fn stop(&self) {
+        let option = self.task_group.lock();
+        if let Some(task_group) = option.as_ref() {
+            task_group.stop();
+        }
+    }
+}
+pub struct TaskGroupGuard {
+    task_group: Arc<Mutex<Option<TaskGroup>>>,
+}
+impl Drop for TaskGroupGuard {
+    fn drop(&mut self) {
+        if let Some(task_group) = self.task_group.lock().take() {
+            task_group.stop();
+        }
+    }
+}