127 lines
3.3 KiB
Rust
127 lines
3.3 KiB
Rust
use crate::net::protocol::{AudioPacket, ControlPacket};
|
|
use anyhow::{Context, Result};
|
|
use std::net::SocketAddr;
|
|
use std::sync::Arc;
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
use tokio::net::{TcpListener, TcpStream, UdpSocket};
|
|
|
|
/// UDP 音频传输
|
|
pub struct AudioSocket {
|
|
socket: Arc<UdpSocket>,
|
|
}
|
|
|
|
impl AudioSocket {
|
|
pub async fn bind() -> Result<Self> {
|
|
let socket = UdpSocket::bind("0.0.0.0:0").await?;
|
|
Ok(Self {
|
|
socket: Arc::new(socket),
|
|
})
|
|
}
|
|
|
|
pub fn local_port(&self) -> Result<u16> {
|
|
Ok(self.socket.local_addr()?.port())
|
|
}
|
|
|
|
pub async fn send_packet(&self, packet: &AudioPacket, target: SocketAddr) -> Result<()> {
|
|
let bytes = postcard::to_allocvec(packet)?;
|
|
self.socket.send_to(&bytes, target).await?;
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn recv_packet(&self, buf: &mut [u8]) -> Result<(AudioPacket, SocketAddr)> {
|
|
let (len, addr) = self.socket.recv_from(buf).await?;
|
|
let packet = postcard::from_bytes(&buf[..len])?;
|
|
Ok((packet, addr))
|
|
}
|
|
|
|
pub async fn punch(&self, target: SocketAddr) -> Result<()> {
|
|
self.socket.send_to(&[0u8; 1], target).await?;
|
|
Ok(())
|
|
}
|
|
|
|
pub fn clone_inner(&self) -> Arc<UdpSocket> {
|
|
self.socket.clone()
|
|
}
|
|
}
|
|
|
|
/// TCP 控制连接
|
|
pub struct ControlConnection {
|
|
stream: TcpStream,
|
|
}
|
|
|
|
impl ControlConnection {
|
|
pub fn new(stream: TcpStream) -> Self {
|
|
Self { stream }
|
|
}
|
|
|
|
pub async fn send_packet(&mut self, packet: &ControlPacket) -> Result<()> {
|
|
let bytes = postcard::to_allocvec(packet)?;
|
|
self.stream.write_all(&bytes).await?;
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn recv_packet(&mut self, buf: &mut [u8]) -> Result<ControlPacket> {
|
|
let len = self.stream.read(buf).await?;
|
|
if len == 0 {
|
|
return Err(anyhow::anyhow!("连接已关闭"));
|
|
}
|
|
let packet = postcard::from_bytes(&buf[..len])?;
|
|
Ok(packet)
|
|
}
|
|
|
|
pub fn split(
|
|
self,
|
|
) -> (
|
|
tokio::net::tcp::OwnedReadHalf,
|
|
tokio::net::tcp::OwnedWriteHalf,
|
|
) {
|
|
self.stream.into_split()
|
|
}
|
|
}
|
|
|
|
/// 主节点网络管理器
|
|
pub struct MasterNetwork {
|
|
listener: TcpListener,
|
|
audio: AudioSocket,
|
|
}
|
|
|
|
impl MasterNetwork {
|
|
pub async fn setup(port: u16) -> Result<Self> {
|
|
let listener = TcpListener::bind(format!("0.0.0.0:{}", port)).await?;
|
|
let audio = AudioSocket::bind().await?;
|
|
Ok(Self { listener, audio })
|
|
}
|
|
|
|
pub async fn accept(&self) -> Result<(ControlConnection, SocketAddr)> {
|
|
let (stream, addr) = self.listener.accept().await?;
|
|
Ok((ControlConnection::new(stream), addr))
|
|
}
|
|
|
|
pub fn audio_socket(&self) -> &AudioSocket {
|
|
&self.audio
|
|
}
|
|
}
|
|
|
|
/// 从节点网络管理器
|
|
pub struct SlaveNetwork {
|
|
control: ControlConnection,
|
|
audio: AudioSocket,
|
|
}
|
|
|
|
impl SlaveNetwork {
|
|
pub async fn connect(master_addr: SocketAddr) -> Result<Self> {
|
|
let stream = TcpStream::connect(master_addr)
|
|
.await
|
|
.context(format!("无法连接到主节点 TCP 地址: {}", master_addr))?;
|
|
let audio = AudioSocket::bind().await?;
|
|
Ok(Self {
|
|
control: ControlConnection::new(stream),
|
|
audio,
|
|
})
|
|
}
|
|
|
|
pub fn split(self) -> (ControlConnection, AudioSocket) {
|
|
(self.control, self.audio)
|
|
}
|
|
}
|