627 lines
23 KiB
Rust
627 lines
23 KiB
Rust
use crate::error::Result;
|
||
use crate::metadata::RbitFetcher;
|
||
use crate::protocol::{DhtMessage, DhtArgs, DhtResponse};
|
||
use crate::scheduler::MetadataScheduler;
|
||
use crate::types::{DHTOptions, TorrentInfo};
|
||
use crate::sharded::{ShardedBloom, ShardedNodeQueue, NodeTuple};
|
||
use rand::Rng;
|
||
use ahash::AHasher;
|
||
use std::hash::{Hash, Hasher};
|
||
use std::net::{IpAddr, SocketAddr};
|
||
use std::sync::{Arc, RwLock};
|
||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||
use std::time::Duration;
|
||
use tokio::net::UdpSocket;
|
||
use tokio::sync::{mpsc, Semaphore};
|
||
use socket2::{Socket, Domain, Type, Protocol};
|
||
use std::pin::Pin;
|
||
use std::future::Future;
|
||
|
||
const BOOTSTRAP_NODES: &[&str] = &[
|
||
"router.bittorrent.com:6881",
|
||
"dht.transmissionbt.com:6881",
|
||
"router.utorrent.com:6881",
|
||
"dht.aelitis.com:6881",
|
||
];
|
||
|
||
// 类型定义
|
||
pub type BoxedBoolFuture = Pin<Box<dyn Future<Output = bool> + Send>>;
|
||
pub type MetadataFetchCallback = Arc<dyn Fn(String) -> BoxedBoolFuture + Send + Sync>;
|
||
|
||
// Hash 发现事件
|
||
/// DHT Server 发现 hash 后发送此事件，由独立的 MetadataScheduler 处理
|
||
#[derive(Debug, Clone)]
|
||
pub struct HashDiscovered {
|
||
pub info_hash: String,
|
||
pub peer_addr: SocketAddr,
|
||
pub discovered_at: std::time::Instant,
|
||
}
|
||
|
||
// ---------------------------------------------------------------
|
||
|
||
type TorrentCallback = Arc<dyn Fn(TorrentInfo) + Send + Sync>;
|
||
type FilterCallback = Arc<dyn Fn(&str) -> bool + Send + Sync>;
|
||
type DuplicateCallback = Arc<dyn Fn(&str) + Send + Sync>;
|
||
|
||
#[derive(Clone)]
|
||
pub struct DHTServer {
|
||
#[allow(dead_code)]
|
||
options: DHTOptions,
|
||
node_id: Vec<u8>,
|
||
socket: Arc<UdpSocket>,
|
||
token_secret: Vec<u8>,
|
||
|
||
// 这些回调现在与 MetadataScheduler 共享
|
||
callback: Arc<RwLock<Option<TorrentCallback>>>,
|
||
filter: Arc<RwLock<Option<FilterCallback>>>,
|
||
on_duplicate: Arc<RwLock<Option<DuplicateCallback>>>,
|
||
on_metadata_fetch: Arc<RwLock<Option<MetadataFetchCallback>>>,
|
||
|
||
// 使用分片锁，大幅减少竞争
|
||
node_queue: Arc<ShardedNodeQueue>,
|
||
seen_hashes: Arc<ShardedBloom>,
|
||
|
||
// 发送 hash 发现事件
|
||
hash_tx: mpsc::Sender<HashDiscovered>,
|
||
|
||
// Metadata 队列长度（用于自适应爬取速度）
|
||
metadata_queue_len: Arc<AtomicUsize>,
|
||
max_metadata_queue_size: usize,
|
||
}
|
||
|
||
impl DHTServer {
|
||
pub async fn new(options: DHTOptions) -> Result<Self> {
|
||
let socket = Socket::new(Domain::IPV4, Type::DGRAM, Some(Protocol::UDP))?;
|
||
#[cfg(not(windows))]
|
||
{ let _ = socket.set_reuse_port(true); }
|
||
let _ = socket.set_reuse_address(true);
|
||
socket.set_nonblocking(true)?;
|
||
|
||
// 增加网络缓冲区以应对高QPS
|
||
let _ = socket.set_recv_buffer_size(32 * 1024 * 1024); // 32MB（原16MB）
|
||
let _ = socket.set_send_buffer_size(8 * 1024 * 1024); // 8MB（原4MB）
|
||
|
||
let addr: SocketAddr = format!("0.0.0.0:{}", options.port).parse().unwrap();
|
||
socket.bind(&addr.into())?;
|
||
let socket = UdpSocket::from_std(socket.into())?;
|
||
|
||
let node_id = generate_random_id();
|
||
let mut rng = rand::thread_rng();
|
||
let token_secret: Vec<u8> = (0..10).map(|_| rng.gen()).collect();
|
||
|
||
// 使用分片队列和分片布隆过滤器
|
||
// 队列容量：100000 个节点（扩容以适应 DHT 网络裂变速度）
|
||
let node_queue = ShardedNodeQueue::new(100000);
|
||
|
||
// 布隆过滤器：预期500万元素，0.1%误判率
|
||
// 内存使用：约 90MB（32分片 × 2.8MB）
|
||
let bloom = ShardedBloom::new_for_fp_rate(5_000_000, 0.001);
|
||
|
||
// -----------------------------------------------------------
|
||
// 内部初始化 MetadataScheduler
|
||
// -----------------------------------------------------------
|
||
let (hash_tx, hash_rx) = mpsc::channel::<HashDiscovered>(10000);
|
||
|
||
let fetcher = Arc::new(RbitFetcher::new(options.metadata_timeout));
|
||
|
||
// 创建共享的回调状态
|
||
let callback = Arc::new(RwLock::new(None));
|
||
let on_metadata_fetch = Arc::new(RwLock::new(None));
|
||
|
||
// 创建共享的队列长度计数器
|
||
let metadata_queue_len = Arc::new(AtomicUsize::new(0));
|
||
|
||
let scheduler = MetadataScheduler::new(
|
||
hash_rx,
|
||
fetcher,
|
||
options.max_metadata_queue_size,
|
||
options.max_metadata_worker_count,
|
||
callback.clone(),
|
||
on_metadata_fetch.clone(),
|
||
metadata_queue_len.clone(),
|
||
);
|
||
|
||
// 启动 Scheduler
|
||
tokio::spawn(async move {
|
||
scheduler.run().await;
|
||
});
|
||
|
||
let server = Self {
|
||
options: options.clone(),
|
||
node_id: node_id.clone(),
|
||
socket: Arc::new(socket),
|
||
token_secret,
|
||
callback,
|
||
on_metadata_fetch,
|
||
node_queue: Arc::new(node_queue),
|
||
seen_hashes: Arc::new(bloom),
|
||
filter: Arc::new(RwLock::new(None)),
|
||
on_duplicate: Arc::new(RwLock::new(None)),
|
||
hash_tx,
|
||
metadata_queue_len,
|
||
max_metadata_queue_size: options.max_metadata_queue_size,
|
||
};
|
||
|
||
Ok(server)
|
||
}
|
||
|
||
pub fn local_addr(&self) -> Result<SocketAddr> {
|
||
Ok(self.socket.local_addr()?)
|
||
}
|
||
|
||
/// 设置元数据获取前的检查回调
|
||
///
|
||
/// 此回调在发现新的 info_hash 后，但在实际连接对等端获取元数据之前执行。
|
||
/// 你可以在这里进行去重检查（如查询数据库），返回 `true` 表示继续获取，`false` 表示跳过。
|
||
///
|
||
/// # 注意事项
|
||
/// - 回调是在 `MetadataScheduler` 的 Worker 线程中异步执行的（通过 `.await`）。
|
||
/// - 支持耗时操作（如数据库查询），但请注意 Worker 数量限制（默认 500）。
|
||
/// - 如果回调执行过慢，可能会导致任务队列堆积。
|
||
///
|
||
/// # 示例
|
||
/// ```rust,ignore
|
||
/// server.on_metadata_fetch(|hash| async move {
|
||
/// // 检查数据库是否存在
|
||
/// // let exists = db.has(hash).await;
|
||
/// // !exists
|
||
/// true
|
||
/// });
|
||
/// ```
|
||
pub fn on_metadata_fetch<F, Fut>(&self, callback: F)
|
||
where
|
||
F: Fn(String) -> Fut + Send + Sync + 'static,
|
||
Fut: Future<Output = bool> + Send + 'static,
|
||
{
|
||
*self.on_metadata_fetch.write().unwrap() = Some(Arc::new(move |hash| {
|
||
Box::pin(callback(hash))
|
||
}));
|
||
}
|
||
|
||
/// 设置成功获取到种子信息的回调
|
||
///
|
||
/// 当成功从对等端下载并解析出种子元数据（Metadata）后调用。
|
||
///
|
||
/// # 注意事项
|
||
/// - 此回调是在 Worker 线程中同步执行的。
|
||
/// - 如果包含耗时操作（如写入大量数据或复杂计算），**必须**在回调内部手动使用 `tokio::spawn`。
|
||
/// - 否则会阻塞当前的元数据获取 Worker，降低系统吞吐量。
|
||
///
|
||
/// # 示例
|
||
/// ```rust,ignore
|
||
/// server.on_torrent(|info| {
|
||
/// // 简单操作可以直接做
|
||
/// println!("Got torrent: {}", info.name);
|
||
///
|
||
/// // 耗时操作建议 spawn
|
||
/// tokio::spawn(async move {
|
||
/// save_to_db(info).await;
|
||
/// });
|
||
/// });
|
||
/// ```
|
||
pub fn on_torrent<F>(&self, callback: F) where F: Fn(TorrentInfo) + Send + Sync + 'static {
|
||
*self.callback.write().unwrap() = Some(Arc::new(callback));
|
||
}
|
||
|
||
/// 设置 Hash 过滤器
|
||
///
|
||
/// 在处理 `announce_peer` 消息时，用于快速判断是否应该处理该 Hash。
|
||
/// 这通常用于布隆过滤器之前的黑名单或白名单机制。
|
||
///
|
||
/// # 注意事项
|
||
/// - 此回调是在 UDP 处理线程中**同步执行**的。
|
||
/// - **绝对禁止**执行任何耗时操作（如 IO、数据库查询、锁等待）。
|
||
/// - 任何延迟都会直接阻塞网络包的接收，导致丢包。
|
||
/// - 应仅进行纯内存的快速判断。
|
||
pub fn set_filter<F>(&self, filter: F) where F: Fn(&str) -> bool + Send + Sync + 'static {
|
||
*self.filter.write().unwrap() = Some(Arc::new(filter));
|
||
}
|
||
|
||
/// 设置重复 Hash 发现的回调
|
||
///
|
||
/// 当接收到的 Hash 已经被布隆过滤器标记为“已存在”时调用。
|
||
///
|
||
/// # 注意事项
|
||
/// - 库内部已经自动为每次调用包裹了 `tokio::spawn`。
|
||
/// - 因此你可以放心地在回调中执行耗时操作（如数据库记录），而不用担心阻塞 UDP 线程。
|
||
/// - 虽然内部有 spawn，但频繁触发仍会产生大量任务，请注意资源控制。
|
||
pub fn on_duplicate<F>(&self, callback: F) where F: Fn(&str) + Send + Sync + 'static {
|
||
*self.on_duplicate.write().unwrap() = Some(Arc::new(callback));
|
||
}
|
||
|
||
pub fn get_seen_count(&self) -> usize {
|
||
// 分片布隆过滤器的位数统计
|
||
self.seen_hashes.number_of_bits() as usize
|
||
}
|
||
|
||
pub fn get_node_pool_size(&self) -> usize {
|
||
self.node_queue.len()
|
||
}
|
||
|
||
pub async fn start(&self) -> Result<()> {
|
||
|
||
self.start_receiver();
|
||
self.bootstrap().await;
|
||
|
||
let server = self.clone();
|
||
|
||
tokio::spawn(async move {
|
||
let semaphore = Arc::new(Semaphore::new(2000));
|
||
let mut loop_tick = 0;
|
||
|
||
loop {
|
||
// 自适应爬取速度：根据 Metadata 队列负载调整爬取策略
|
||
let queue_len = server.metadata_queue_len.load(Ordering::Relaxed);
|
||
let queue_pressure = queue_len as f64 / server.max_metadata_queue_size as f64;
|
||
|
||
// 动态计算批次大小和休眠时间
|
||
let (batch_size, sleep_duration) = if queue_pressure < 0.5 {
|
||
// 🟢 绿区：队列空闲，全速爬取
|
||
(200, Duration::from_millis(10))
|
||
} else if queue_pressure < 0.8 {
|
||
// 🟡 黄区：队列有压力，适度减速
|
||
(200, Duration::from_millis(20))
|
||
} else if queue_pressure < 0.95 {
|
||
// 🟠 橙区：队列高压，大幅减速
|
||
(20, Duration::from_millis(500))
|
||
} else {
|
||
// 🔴 红区：队列爆满，暂停主动爬取
|
||
(0, Duration::from_millis(1000))
|
||
};
|
||
|
||
let nodes_batch = {
|
||
if server.node_queue.is_empty() || batch_size == 0 {
|
||
None
|
||
} else {
|
||
Some(server.node_queue.pop_batch(batch_size))
|
||
}
|
||
};
|
||
|
||
loop_tick += 1;
|
||
if nodes_batch.is_none() || loop_tick % 50 == 0 {
|
||
server.bootstrap().await;
|
||
if nodes_batch.is_none() {
|
||
tokio::time::sleep(sleep_duration).await;
|
||
continue;
|
||
}
|
||
}
|
||
|
||
if let Some(nodes) = nodes_batch {
|
||
for node in nodes {
|
||
let permit = semaphore.clone().acquire_owned().await.unwrap();
|
||
let server_clone = server.clone();
|
||
tokio::spawn(async move {
|
||
let neighbor_id = generate_neighbor_target(&node.id, &server_clone.node_id);
|
||
let random_target = generate_random_id();
|
||
let _ = server_clone.send_find_node(node.addr, &random_target, &neighbor_id).await;
|
||
drop(permit);
|
||
});
|
||
}
|
||
}
|
||
|
||
tokio::time::sleep(sleep_duration).await;
|
||
}
|
||
});
|
||
|
||
std::future::pending::<()>().await;
|
||
Ok(())
|
||
}
|
||
|
||
fn start_receiver(&self) {
|
||
let socket = self.socket.clone();
|
||
let server = self.clone();
|
||
|
||
let num_workers = std::thread::available_parallelism()
|
||
.map(|n| n.get())
|
||
.unwrap_or(8);
|
||
|
||
let queue_size = 5000;
|
||
|
||
let mut senders = Vec::with_capacity(num_workers);
|
||
for _ in 0..num_workers {
|
||
let (tx, mut rx) = mpsc::channel::<(Vec<u8>, SocketAddr)>(queue_size);
|
||
senders.push(tx);
|
||
|
||
let server_clone = server.clone();
|
||
|
||
tokio::spawn(async move {
|
||
while let Some((data, addr)) = rx.recv().await {
|
||
let _ = server_clone.handle_message(&data, addr).await;
|
||
}
|
||
});
|
||
}
|
||
|
||
tokio::spawn(async move {
|
||
let mut buf = [0u8; 65536];
|
||
let mut next_worker_idx = 0;
|
||
|
||
loop {
|
||
match socket.recv_from(&mut buf).await {
|
||
Ok((size, addr)) => {
|
||
// 🛡️ 安全检查1：拒绝异常大的包（DHT 消息通常 < 2KB）
|
||
if size > 8192 {
|
||
#[cfg(debug_assertions)]
|
||
log::trace!("⚠️ 拒绝异常大的 UDP 包: {} 字节 from {}", size, addr);
|
||
continue;
|
||
}
|
||
|
||
// 🛡️ 安全检查2：快速检查是否是有效的 Bencode 字典
|
||
// DHT KRPC 消息（BEP-5）必须是字典，首字符必须是 'd'
|
||
if size == 0 || buf[0] != b'd' {
|
||
continue;
|
||
}
|
||
|
||
let data = buf[..size].to_vec();
|
||
|
||
let tx = &senders[next_worker_idx];
|
||
next_worker_idx = (next_worker_idx + 1) % num_workers;
|
||
|
||
match tx.try_send((data, addr)) {
|
||
Ok(_) => {},
|
||
Err(mpsc::error::TrySendError::Full(_)) => {
|
||
#[cfg(debug_assertions)]
|
||
log::trace!("UDP worker queue full, dropping packet");
|
||
},
|
||
Err(_) => { break; }
|
||
}
|
||
}
|
||
Err(_e) => {
|
||
tokio::time::sleep(Duration::from_millis(1)).await;
|
||
}
|
||
}
|
||
}
|
||
});
|
||
}
|
||
|
||
async fn handle_message(&self, data: &[u8], addr: SocketAddr) -> Result<()> {
|
||
let msg: DhtMessage = match serde_bencode::from_bytes(data) {
|
||
Ok(m) => m,
|
||
Err(_) => return Ok(()),
|
||
};
|
||
|
||
match msg.y.as_str() {
|
||
"q" => {
|
||
if let Some(q_type) = &msg.q {
|
||
self.handle_query(&msg, q_type.as_bytes(), addr).await?;
|
||
}
|
||
}
|
||
"r" => {
|
||
if let Some(response) = &msg.r {
|
||
self.handle_response(response).await?;
|
||
}
|
||
}
|
||
_ => {}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
async fn handle_query(&self, msg: &DhtMessage, query_type: &[u8], addr: SocketAddr) -> Result<()> {
|
||
let args = match &msg.a {
|
||
Some(a) => a,
|
||
None => return Ok(()),
|
||
};
|
||
|
||
let transaction_id = &msg.t;
|
||
let sender_id: Option<&[u8]> = args.id.as_deref().map(|v| v.as_slice());
|
||
let target_id_fallback: Option<&[u8]> = args.target.as_deref()
|
||
.or(args.info_hash.as_deref())
|
||
.map(|v| v.as_slice());
|
||
|
||
let q_str = std::str::from_utf8(query_type).unwrap_or("");
|
||
|
||
if q_str == "announce_peer" {
|
||
self.handle_announce_peer(args, addr).await?;
|
||
}
|
||
|
||
self.send_response(transaction_id, addr, q_str, sender_id, target_id_fallback).await?;
|
||
Ok(())
|
||
}
|
||
|
||
async fn handle_announce_peer(&self, args: &DhtArgs, addr: SocketAddr) -> Result<()> {
|
||
if let Some(token) = &args.token {
|
||
if !self.validate_token(token, addr) { return Ok(()); }
|
||
} else {
|
||
return Ok(());
|
||
}
|
||
|
||
if let Some(info_hash) = &args.info_hash {
|
||
let info_hash_arr: [u8; 20] = match info_hash.as_ref().try_into() {
|
||
Ok(arr) => arr, Err(_) => return Ok(()),
|
||
};
|
||
let hash_hex = hex::encode(info_hash_arr);
|
||
|
||
// 使用分片布隆过滤器进行高效去重
|
||
let is_duplicate = self.seen_hashes.check_and_set(&info_hash_arr);
|
||
|
||
if is_duplicate {
|
||
let dup_cb = self.on_duplicate.read().unwrap().clone();
|
||
if let Some(cb) = dup_cb {
|
||
let hash_hex_clone = hash_hex.clone();
|
||
tokio::spawn(async move {
|
||
cb(&hash_hex_clone);
|
||
});
|
||
}
|
||
return Ok(());
|
||
}
|
||
|
||
let filter_cb = self.filter.read().unwrap().clone();
|
||
if let Some(f) = filter_cb {
|
||
if !f(&hash_hex) { return Ok(()); }
|
||
}
|
||
|
||
#[cfg(debug_assertions)]
|
||
log::debug!("🔥 新 Hash: {} 来自 {}", hash_hex, addr);
|
||
|
||
// 解耦：发送 hash 发现事件
|
||
let port = if let Some(implied) = args.implied_port {
|
||
if implied != 0 { addr.port() } else { args.port.unwrap_or(0) }
|
||
} else {
|
||
args.port.unwrap_or(addr.port())
|
||
};
|
||
|
||
if port > 0 {
|
||
let event = HashDiscovered {
|
||
info_hash: hash_hex,
|
||
peer_addr: SocketAddr::new(addr.ip(), port),
|
||
discovered_at: std::time::Instant::now(),
|
||
};
|
||
|
||
// 使用 try_send，队列满时直接丢弃（背压）
|
||
if let Err(_) = self.hash_tx.try_send(event) {
|
||
#[cfg(debug_assertions)]
|
||
log::trace!("⚠️ Hash 队列满，丢弃 hash");
|
||
}
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
async fn handle_response(&self, response: &DhtResponse) -> Result<()> {
|
||
if let Some(nodes_bytes) = &response.nodes {
|
||
self.process_compact_nodes(nodes_bytes);
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
fn process_compact_nodes(&self, nodes_bytes: &[u8]) {
|
||
if nodes_bytes.len() % 26 != 0 { return; }
|
||
|
||
// 使用分片队列，直接并发插入（无锁竞争）
|
||
for chunk in nodes_bytes.chunks(26) {
|
||
let id = chunk[0..20].to_vec();
|
||
let port = u16::from_be_bytes([chunk[24], chunk[25]]);
|
||
|
||
let ip = std::net::Ipv4Addr::new(chunk[20], chunk[21], chunk[22], chunk[23]);
|
||
let addr = SocketAddr::new(std::net::IpAddr::V4(ip), port);
|
||
|
||
self.node_queue.push(NodeTuple { id, addr });
|
||
}
|
||
}
|
||
|
||
async fn send_response(
|
||
&self,
|
||
tid: &[u8],
|
||
addr: SocketAddr,
|
||
query_type: &str,
|
||
sender_id: Option<&[u8]>,
|
||
target_id_fallback: Option<&[u8]>,
|
||
) -> Result<()> {
|
||
let mut r_dict = std::collections::HashMap::new();
|
||
|
||
let reference_id = sender_id.or(target_id_fallback);
|
||
let my_id = if let Some(target) = reference_id {
|
||
generate_neighbor_target(target, &self.node_id)
|
||
} else {
|
||
self.node_id.clone()
|
||
};
|
||
|
||
r_dict.insert(b"id".to_vec(), serde_bencode::value::Value::Bytes(my_id));
|
||
let token = self.generate_token(addr);
|
||
r_dict.insert(b"token".to_vec(), serde_bencode::value::Value::Bytes(token));
|
||
|
||
if query_type == "get_peers" || query_type == "find_node" {
|
||
// 使用分片队列获取随机节点（无锁竞争）
|
||
let nodes = self.node_queue.get_random_nodes(8);
|
||
|
||
let mut nodes_data = Vec::new();
|
||
for node in nodes {
|
||
nodes_data.extend_from_slice(&node.id);
|
||
match node.addr.ip() {
|
||
IpAddr::V4(ip) => nodes_data.extend_from_slice(&ip.octets()),
|
||
_ => continue,
|
||
}
|
||
nodes_data.extend_from_slice(&node.addr.port().to_be_bytes());
|
||
}
|
||
|
||
r_dict.insert(b"nodes".to_vec(), serde_bencode::value::Value::Bytes(nodes_data));
|
||
}
|
||
|
||
let mut response: std::collections::HashMap<String, serde_bencode::value::Value> = std::collections::HashMap::new();
|
||
response.insert("t".to_string(), serde_bencode::value::Value::Bytes(tid.to_vec()));
|
||
response.insert("y".to_string(), serde_bencode::value::Value::Bytes(b"r".to_vec()));
|
||
response.insert("r".to_string(), serde_bencode::value::Value::Dict(r_dict));
|
||
|
||
if let Ok(encoded) = serde_bencode::to_bytes(&response) {
|
||
let _ = self.socket.send_to(&encoded, addr).await;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
async fn bootstrap(&self) {
|
||
let target = generate_random_id();
|
||
for node in BOOTSTRAP_NODES {
|
||
match tokio::net::lookup_host(node).await {
|
||
Ok(addrs) => {
|
||
for addr in addrs {
|
||
if addr.is_ipv6() { continue; }
|
||
let _ = self.send_find_node(addr, &target, &self.node_id).await;
|
||
}
|
||
}
|
||
Err(_) => {}
|
||
}
|
||
}
|
||
}
|
||
|
||
async fn send_find_node(&self, addr: SocketAddr, target: &[u8], sender_id: &[u8]) -> Result<()> {
|
||
let mut args = std::collections::HashMap::new();
|
||
args.insert(b"id".to_vec(), serde_bencode::value::Value::Bytes(sender_id.to_vec()));
|
||
args.insert(b"target".to_vec(), serde_bencode::value::Value::Bytes(target.to_vec()));
|
||
|
||
let mut msg: std::collections::HashMap<String, serde_bencode::value::Value> = std::collections::HashMap::new();
|
||
msg.insert("t".to_string(), serde_bencode::value::Value::Bytes(vec![0, 1]));
|
||
msg.insert("y".to_string(), serde_bencode::value::Value::Bytes(b"q".to_vec()));
|
||
msg.insert("q".to_string(), serde_bencode::value::Value::Bytes(b"find_node".to_vec()));
|
||
msg.insert("a".to_string(), serde_bencode::value::Value::Dict(args));
|
||
|
||
if let Ok(encoded) = serde_bencode::to_bytes(&msg) {
|
||
let _ = self.socket.send_to(&encoded, addr).await;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
fn generate_token(&self, addr: SocketAddr) -> Vec<u8> {
|
||
|
||
let mut hasher = AHasher::default();
|
||
|
||
// Hash IP地址
|
||
match addr.ip() {
|
||
IpAddr::V4(ip) => ip.octets().hash(&mut hasher),
|
||
IpAddr::V6(ip) => ip.octets().hash(&mut hasher),
|
||
}
|
||
|
||
// Hash 密钥
|
||
self.token_secret.hash(&mut hasher);
|
||
|
||
// 返回 8 字节 token
|
||
let hash = hasher.finish();
|
||
hash.to_le_bytes().to_vec()
|
||
}
|
||
|
||
fn validate_token(&self, token: &[u8], addr: SocketAddr) -> bool {
|
||
if token.len() != 8 {
|
||
return false;
|
||
}
|
||
let expected = self.generate_token(addr);
|
||
token == expected.as_slice()
|
||
}
|
||
}
|
||
|
||
fn generate_random_id() -> Vec<u8> {
|
||
let mut rng = rand::thread_rng();
|
||
(0..20).map(|_| rng.gen()).collect()
|
||
}
|
||
|
||
fn generate_neighbor_target(remote_id: &[u8], local_id: &[u8]) -> Vec<u8> {
|
||
let mut id = Vec::with_capacity(20);
|
||
let prefix_len = std::cmp::min(remote_id.len(), 6);
|
||
id.extend_from_slice(&remote_id[..prefix_len]);
|
||
if local_id.len() > prefix_len {
|
||
id.extend_from_slice(&local_id[prefix_len..]);
|
||
} else {
|
||
while id.len() < 20 {
|
||
id.push(rand::random());
|
||
}
|
||
}
|
||
id
|
||
}
|