482 lines
16 KiB
Rust
482 lines
16 KiB
Rust
//! P2P 连接模块
|
||
|
||
use anyhow::Result;
|
||
use bytes::Bytes;
|
||
use easyremote_common::protocol::{FrameData, FrameFormat, InputEvent};
|
||
use std::net::SocketAddr;
|
||
use std::sync::Arc;
|
||
use tokio::net::UdpSocket;
|
||
use tokio::sync::{mpsc, RwLock};
|
||
|
||
/// 连接状态
|
||
#[derive(Debug, Clone, PartialEq)]
|
||
pub enum ConnectionState {
|
||
/// 断开连接
|
||
Disconnected,
|
||
/// 正在连接
|
||
Connecting,
|
||
/// 已连接
|
||
Connected,
|
||
/// 连接失败
|
||
Failed(String),
|
||
}
|
||
|
||
/// P2P 连接管理器
|
||
pub struct P2PConnection {
|
||
state: Arc<RwLock<ConnectionState>>,
|
||
local_addr: Option<SocketAddr>,
|
||
remote_addr: Option<SocketAddr>,
|
||
socket: Option<Arc<UdpSocket>>,
|
||
frame_tx: Option<mpsc::Sender<FrameData>>,
|
||
input_tx: Option<mpsc::Sender<InputEvent>>,
|
||
}
|
||
|
||
impl P2PConnection {
|
||
pub fn new() -> Self {
|
||
Self {
|
||
state: Arc::new(RwLock::new(ConnectionState::Disconnected)),
|
||
local_addr: None,
|
||
remote_addr: None,
|
||
socket: None,
|
||
frame_tx: None,
|
||
input_tx: None,
|
||
}
|
||
}
|
||
|
||
/// 获取连接状态
|
||
pub async fn state(&self) -> ConnectionState {
|
||
self.state.read().await.clone()
|
||
}
|
||
|
||
/// 初始化本地套接字
|
||
pub async fn init_socket(&mut self) -> Result<SocketAddr> {
|
||
let socket = UdpSocket::bind("0.0.0.0:0").await?;
|
||
let local_addr = socket.local_addr()?;
|
||
self.local_addr = Some(local_addr);
|
||
self.socket = Some(Arc::new(socket));
|
||
Ok(local_addr)
|
||
}
|
||
|
||
/// 尝试连接到远程地址
|
||
pub async fn connect(&mut self, remote_addr: SocketAddr) -> Result<()> {
|
||
*self.state.write().await = ConnectionState::Connecting;
|
||
|
||
if let Some(socket) = &self.socket {
|
||
socket.connect(remote_addr).await?;
|
||
self.remote_addr = Some(remote_addr);
|
||
*self.state.write().await = ConnectionState::Connected;
|
||
} else {
|
||
*self.state.write().await = ConnectionState::Failed("Socket not initialized".into());
|
||
anyhow::bail!("Socket not initialized");
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// 发送帧数据
|
||
pub async fn send_frame(&self, frame: &FrameData) -> Result<()> {
|
||
if let Some(socket) = &self.socket {
|
||
let data = bincode::serialize(frame)?;
|
||
// 如果数据太大，需要分片发送
|
||
if data.len() > 65000 {
|
||
self.send_fragmented(socket, &data).await?;
|
||
} else {
|
||
socket.send(&data).await?;
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// 发送输入事件
|
||
pub async fn send_input(&self, event: &InputEvent) -> Result<()> {
|
||
if let Some(socket) = &self.socket {
|
||
let data = bincode::serialize(event)?;
|
||
socket.send(&data).await?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// 分片发送大数据
|
||
async fn send_fragmented(&self, socket: &UdpSocket, data: &[u8]) -> Result<()> {
|
||
const MAX_FRAGMENT_SIZE: usize = 60000;
|
||
let total_fragments = (data.len() + MAX_FRAGMENT_SIZE - 1) / MAX_FRAGMENT_SIZE;
|
||
|
||
for (i, chunk) in data.chunks(MAX_FRAGMENT_SIZE).enumerate() {
|
||
let fragment = FragmentHeader {
|
||
fragment_id: i as u16,
|
||
total_fragments: total_fragments as u16,
|
||
data: chunk.to_vec(),
|
||
};
|
||
let fragment_data = bincode::serialize(&fragment)?;
|
||
socket.send(&fragment_data).await?;
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// 开始接收数据
|
||
pub async fn start_receiving(
|
||
&self,
|
||
on_frame: impl Fn(FrameData) + Send + 'static,
|
||
on_input: impl Fn(InputEvent) + Send + 'static,
|
||
) -> Result<()> {
|
||
let socket = self.socket.clone().ok_or_else(|| anyhow::anyhow!("Socket not initialized"))?;
|
||
let state = self.state.clone();
|
||
|
||
tokio::spawn(async move {
|
||
let mut buf = vec![0u8; 100000];
|
||
let mut fragment_buffer: Vec<Option<Vec<u8>>> = Vec::new();
|
||
let mut expected_fragments = 0u16;
|
||
|
||
loop {
|
||
if *state.read().await != ConnectionState::Connected {
|
||
break;
|
||
}
|
||
|
||
match socket.recv(&mut buf).await {
|
||
Ok(len) => {
|
||
let data = &buf[..len];
|
||
|
||
// 尝试解析为分片
|
||
if let Ok(fragment) = bincode::deserialize::<FragmentHeader>(data) {
|
||
if fragment.total_fragments > 1 {
|
||
// 处理分片
|
||
if fragment_buffer.is_empty() {
|
||
fragment_buffer = vec![None; fragment.total_fragments as usize];
|
||
expected_fragments = fragment.total_fragments;
|
||
}
|
||
|
||
if fragment.fragment_id < expected_fragments {
|
||
fragment_buffer[fragment.fragment_id as usize] =
|
||
Some(fragment.data);
|
||
}
|
||
|
||
// 检查是否收到所有分片
|
||
if fragment_buffer.iter().all(|f| f.is_some()) {
|
||
let complete_data: Vec<u8> = fragment_buffer
|
||
.iter()
|
||
.filter_map(|f| f.clone())
|
||
.flatten()
|
||
.collect();
|
||
|
||
if let Ok(frame) = bincode::deserialize::<FrameData>(&complete_data) {
|
||
on_frame(frame);
|
||
}
|
||
|
||
fragment_buffer.clear();
|
||
}
|
||
continue;
|
||
}
|
||
}
|
||
|
||
// 尝试解析为帧数据
|
||
if let Ok(frame) = bincode::deserialize::<FrameData>(data) {
|
||
on_frame(frame);
|
||
continue;
|
||
}
|
||
|
||
// 尝试解析为输入事件
|
||
if let Ok(input) = bincode::deserialize::<InputEvent>(data) {
|
||
on_input(input);
|
||
}
|
||
}
|
||
Err(e) => {
|
||
tracing::warn!("Receive error: {}", e);
|
||
}
|
||
}
|
||
}
|
||
});
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// 断开连接
|
||
pub async fn disconnect(&mut self) {
|
||
*self.state.write().await = ConnectionState::Disconnected;
|
||
self.socket = None;
|
||
self.remote_addr = None;
|
||
}
|
||
}
|
||
|
||
/// 分片头
|
||
#[derive(Debug, serde::Serialize, serde::Deserialize)]
|
||
struct FragmentHeader {
|
||
fragment_id: u16,
|
||
total_fragments: u16,
|
||
data: Vec<u8>,
|
||
}
|
||
|
||
/// ICE 服务器配置
|
||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||
pub struct IceServersConfig {
|
||
pub stun_servers: Vec<String>,
|
||
pub turn_server: Option<TurnConfig>,
|
||
}
|
||
|
||
/// TURN 服务器配置
|
||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||
pub struct TurnConfig {
|
||
pub url: String,
|
||
pub username: String,
|
||
pub credential: String,
|
||
}
|
||
|
||
/// NAT穿透辅助
|
||
pub struct NatTraversal;
|
||
|
||
impl NatTraversal {
|
||
/// 获取所有本地 IP 地址
|
||
pub fn get_local_ips() -> Vec<String> {
|
||
let mut ips = Vec::new();
|
||
|
||
// 尝试使用 socket 连接来获取本地 IP
|
||
if let Ok(socket) = std::net::UdpSocket::bind("0.0.0.0:0") {
|
||
// 连接到公网地址来获取本机实际使用的 IP
|
||
if let Ok(()) = socket.connect("8.8.8.8:80") {
|
||
if let Ok(addr) = socket.local_addr() {
|
||
let ip = addr.ip().to_string();
|
||
if !ip.starts_with("127.") && ip != "0.0.0.0" {
|
||
ips.push(ip);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// 如果没有获取到，添加 localhost
|
||
if ips.is_empty() {
|
||
ips.push("127.0.0.1".to_string());
|
||
}
|
||
|
||
ips
|
||
}
|
||
|
||
/// 从服务器获取 ICE 服务器配置
|
||
pub async fn fetch_ice_servers(server_url: &str) -> Result<IceServersConfig> {
|
||
// 将 ws:// 或 wss:// 转换为 http:// 或 https://
|
||
let http_url = server_url
|
||
.replace("ws://", "http://")
|
||
.replace("wss://", "https://");
|
||
|
||
let url = format!("{}/api/ice-servers", http_url.trim_end_matches('/'));
|
||
|
||
let response = reqwest::get(&url).await?;
|
||
let config: IceServersConfig = response.json().await?;
|
||
|
||
Ok(config)
|
||
}
|
||
|
||
/// 获取本地候选地址
|
||
pub async fn get_local_candidates() -> Result<Vec<String>> {
|
||
let mut candidates = Vec::new();
|
||
|
||
// 获取所有本地网络接口
|
||
if let Ok(socket) = UdpSocket::bind("0.0.0.0:0").await {
|
||
if let Ok(addr) = socket.local_addr() {
|
||
candidates.push(format!("host {}", addr));
|
||
}
|
||
}
|
||
|
||
Ok(candidates)
|
||
}
|
||
|
||
/// 获取完整的 ICE 候选（包括本地和公网地址）
|
||
pub async fn get_all_candidates(stun_servers: &[String]) -> Result<Vec<IceCandidate>> {
|
||
let mut candidates = Vec::new();
|
||
|
||
// 添加本地地址候选
|
||
let local_candidates = Self::get_local_candidates().await?;
|
||
for candidate in local_candidates {
|
||
candidates.push(IceCandidate::new(candidate));
|
||
}
|
||
|
||
// 使用 STUN 服务器获取公网地址
|
||
for stun_url in stun_servers {
|
||
match Self::get_public_addr_from_url(stun_url).await {
|
||
Ok(addr) => {
|
||
candidates.push(IceCandidate::new(format!("srflx {}", addr)));
|
||
tracing::info!("Got public address from STUN {}: {}", stun_url, addr);
|
||
break; // 成功获取一个公网地址即可
|
||
}
|
||
Err(e) => {
|
||
tracing::warn!("Failed to get public address from {}: {}", stun_url, e);
|
||
}
|
||
}
|
||
}
|
||
|
||
Ok(candidates)
|
||
}
|
||
|
||
/// 从 STUN URL 解析并获取公网地址
|
||
pub async fn get_public_addr_from_url(stun_url: &str) -> Result<SocketAddr> {
|
||
// 解析 STUN URL，格式: stun:host:port 或 stun:host
|
||
let url = stun_url.trim_start_matches("stun:");
|
||
let parts: Vec<&str> = url.split(':').collect();
|
||
|
||
let (host, port) = match parts.len() {
|
||
1 => (parts[0], 3478u16), // 默认 STUN 端口
|
||
2 => (parts[0], parts[1].parse().unwrap_or(3478)),
|
||
_ => anyhow::bail!("Invalid STUN URL format: {}", stun_url),
|
||
};
|
||
|
||
Self::get_public_addr(host, port).await
|
||
}
|
||
|
||
/// 使用 STUN 服务器获取公网地址
|
||
pub async fn get_public_addr(stun_host: &str, stun_port: u16) -> Result<SocketAddr> {
|
||
let socket = UdpSocket::bind("0.0.0.0:0").await?;
|
||
|
||
// 解析 STUN 服务器地址
|
||
let stun_addr = tokio::net::lookup_host(format!("{}:{}", stun_host, stun_port))
|
||
.await?
|
||
.next()
|
||
.ok_or_else(|| anyhow::anyhow!("Cannot resolve STUN server: {}", stun_host))?;
|
||
|
||
// 设置超时
|
||
socket.connect(stun_addr).await?;
|
||
|
||
// 发送 STUN Binding 请求
|
||
let binding_request = create_stun_binding_request();
|
||
socket.send(&binding_request).await?;
|
||
|
||
// 接收响应（带超时）
|
||
let mut buf = [0u8; 1024];
|
||
let recv_result = tokio::time::timeout(
|
||
std::time::Duration::from_secs(5),
|
||
socket.recv(&mut buf)
|
||
).await;
|
||
|
||
match recv_result {
|
||
Ok(Ok(len)) => parse_stun_response(&buf[..len]),
|
||
Ok(Err(e)) => anyhow::bail!("STUN receive error: {}", e),
|
||
Err(_) => anyhow::bail!("STUN request timeout"),
|
||
}
|
||
}
|
||
}
|
||
|
||
/// STUN Magic Cookie (RFC 5389)
|
||
const STUN_MAGIC_COOKIE: u32 = 0x2112A442;
|
||
|
||
/// STUN Attribute Types
|
||
const ATTR_XOR_MAPPED_ADDRESS: u16 = 0x0020;
|
||
|
||
fn create_stun_binding_request() -> Vec<u8> {
|
||
use rand::Rng;
|
||
|
||
let mut request = Vec::with_capacity(20);
|
||
|
||
// Message Type: Binding Request (0x0001)
|
||
request.extend_from_slice(&0x0001u16.to_be_bytes());
|
||
// Message Length: 0 (no attributes)
|
||
request.extend_from_slice(&0x0000u16.to_be_bytes());
|
||
// Magic Cookie (RFC 5389)
|
||
request.extend_from_slice(&STUN_MAGIC_COOKIE.to_be_bytes());
|
||
// Transaction ID (12 bytes random)
|
||
let mut rng = rand::thread_rng();
|
||
let transaction_id: [u8; 12] = rng.gen();
|
||
request.extend_from_slice(&transaction_id);
|
||
|
||
request
|
||
}
|
||
|
||
fn parse_stun_response(data: &[u8]) -> Result<SocketAddr> {
|
||
// 验证最小长度 (20 bytes header)
|
||
if data.len() < 20 {
|
||
anyhow::bail!("Invalid STUN response: too short");
|
||
}
|
||
|
||
// 验证 Magic Cookie
|
||
let magic = u32::from_be_bytes([data[4], data[5], data[6], data[7]]);
|
||
if magic != STUN_MAGIC_COOKIE {
|
||
anyhow::bail!("Invalid STUN magic cookie");
|
||
}
|
||
|
||
// 解析消息长度
|
||
let msg_len = u16::from_be_bytes([data[2], data[3]]) as usize;
|
||
if data.len() < 20 + msg_len {
|
||
anyhow::bail!("Invalid STUN response: truncated");
|
||
}
|
||
|
||
// 遍历属性查找 XOR-MAPPED-ADDRESS
|
||
let mut offset = 20;
|
||
while offset + 4 <= 20 + msg_len {
|
||
let attr_type = u16::from_be_bytes([data[offset], data[offset + 1]]);
|
||
let attr_len = u16::from_be_bytes([data[offset + 2], data[offset + 3]]) as usize;
|
||
|
||
if attr_type == ATTR_XOR_MAPPED_ADDRESS {
|
||
// XOR-MAPPED-ADDRESS 格式:
|
||
// 1 byte: reserved
|
||
// 1 byte: family (0x01=IPv4, 0x02=IPv6)
|
||
// 2 bytes: XOR'd port
|
||
// 4 bytes (IPv4) or 16 bytes (IPv6): XOR'd address
|
||
|
||
if offset + 4 + attr_len > data.len() {
|
||
anyhow::bail!("Invalid XOR-MAPPED-ADDRESS attribute");
|
||
}
|
||
|
||
let family = data[offset + 5];
|
||
|
||
// XOR'd port
|
||
let xor_port = u16::from_be_bytes([data[offset + 6], data[offset + 7]]);
|
||
let port = xor_port ^ ((STUN_MAGIC_COOKIE >> 16) as u16);
|
||
|
||
match family {
|
||
0x01 => {
|
||
// IPv4
|
||
let magic_bytes = STUN_MAGIC_COOKIE.to_be_bytes();
|
||
let ip = std::net::Ipv4Addr::new(
|
||
data[offset + 8] ^ magic_bytes[0],
|
||
data[offset + 9] ^ magic_bytes[1],
|
||
data[offset + 10] ^ magic_bytes[2],
|
||
data[offset + 11] ^ magic_bytes[3],
|
||
);
|
||
return Ok(SocketAddr::new(std::net::IpAddr::V4(ip), port));
|
||
}
|
||
0x02 => {
|
||
// IPv6 (需要 Transaction ID 进行 XOR)
|
||
let magic_bytes = STUN_MAGIC_COOKIE.to_be_bytes();
|
||
let transaction_id = &data[8..20];
|
||
let mut ip_bytes = [0u8; 16];
|
||
|
||
// 前 4 bytes 与 magic cookie XOR
|
||
for i in 0..4 {
|
||
ip_bytes[i] = data[offset + 8 + i] ^ magic_bytes[i];
|
||
}
|
||
// 后 12 bytes 与 transaction ID XOR
|
||
for i in 4..16 {
|
||
ip_bytes[i] = data[offset + 8 + i] ^ transaction_id[i - 4];
|
||
}
|
||
|
||
let ip = std::net::Ipv6Addr::from(ip_bytes);
|
||
return Ok(SocketAddr::new(std::net::IpAddr::V6(ip), port));
|
||
}
|
||
_ => {
|
||
anyhow::bail!("Unknown address family: {}", family);
|
||
}
|
||
}
|
||
}
|
||
|
||
// 移动到下一个属性 (4字节对齐)
|
||
offset += 4 + attr_len + (4 - (attr_len % 4)) % 4;
|
||
}
|
||
|
||
anyhow::bail!("XOR-MAPPED-ADDRESS not found in STUN response")
|
||
}
|
||
|
||
/// ICE 候选
|
||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||
pub struct IceCandidate {
|
||
pub candidate: String,
|
||
pub sdp_mid: Option<String>,
|
||
pub sdp_mline_index: Option<u32>,
|
||
}
|
||
|
||
impl IceCandidate {
|
||
pub fn new(candidate: String) -> Self {
|
||
Self {
|
||
candidate,
|
||
sdp_mid: None,
|
||
sdp_mline_index: Some(0),
|
||
}
|
||
}
|
||
}
|