126 lines
4.0 KiB
Rust
126 lines
4.0 KiB
Rust
use crate::network::{NetworkTransport, NetError};
|
|
use iroh::{Endpoint, EndpointId};
|
|
use iroh::endpoint::Connection;
|
|
use bytes::Bytes;
|
|
use tokio::sync::{mpsc, Mutex};
|
|
use tokio::sync::mpsc::Receiver;
|
|
use std::sync::Arc;
|
|
use std::collections::HashMap;
|
|
use async_trait::async_trait;
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct AudioProtocol {
|
|
incoming_tx: mpsc::Sender<(EndpointId, Bytes)>,
|
|
connections: Arc<Mutex<HashMap<EndpointId, Connection>>>,
|
|
}
|
|
|
|
impl iroh::protocol::ProtocolHandler for AudioProtocol {
|
|
fn accept(
|
|
&self,
|
|
connection: Connection,
|
|
) -> impl std::future::Future<Output = Result<(), iroh::protocol::AcceptError>> + Send {
|
|
let peer_id = connection.remote_id();
|
|
let incoming_tx = self.incoming_tx.clone();
|
|
let connections = self.connections.clone();
|
|
|
|
async move {
|
|
connections.lock().await.insert(peer_id, connection.clone());
|
|
loop {
|
|
match connection.read_datagram().await {
|
|
Ok(bytes) => {
|
|
if incoming_tx.send((peer_id, bytes)).await.is_err() {
|
|
break;
|
|
}
|
|
}
|
|
Err(_) => {
|
|
connections.lock().await.remove(&peer_id);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
}
|
|
|
|
pub struct IrohTransport {
|
|
endpoint: Endpoint,
|
|
connections: Arc<Mutex<HashMap<EndpointId, Connection>>>,
|
|
incoming_tx: mpsc::Sender<(EndpointId, Bytes)>,
|
|
incoming_rx: Mutex<Option<mpsc::Receiver<(EndpointId, Bytes)>>>,
|
|
}
|
|
|
|
impl IrohTransport {
|
|
pub fn new(endpoint: Endpoint) -> (Self, AudioProtocol) {
|
|
let (incoming_tx, incoming_rx) = mpsc::channel(1000);
|
|
let connections = Arc::new(Mutex::new(HashMap::new()));
|
|
|
|
let audio_proto = AudioProtocol {
|
|
incoming_tx: incoming_tx.clone(),
|
|
connections: connections.clone(),
|
|
};
|
|
|
|
let transport = Self {
|
|
endpoint,
|
|
connections,
|
|
incoming_tx,
|
|
incoming_rx: Mutex::new(Some(incoming_rx)),
|
|
};
|
|
|
|
(transport, audio_proto)
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl NetworkTransport for IrohTransport {
|
|
async fn send_datagram(&self, peer_id: EndpointId, data: Bytes) -> Result<(), NetError> {
|
|
let mut conns = self.connections.lock().await;
|
|
let conn = if let Some(conn) = conns.get(&peer_id) {
|
|
conn.clone()
|
|
} else {
|
|
// Establish a new connection.
|
|
// We use the same audio ALPN: b"peerspeak-audio"
|
|
let alpn = b"peerspeak-audio";
|
|
let conn = self.endpoint.connect(peer_id, alpn).await
|
|
.map_err(|e| NetError::Connection(e.to_string()))?;
|
|
|
|
conns.insert(peer_id, conn.clone());
|
|
|
|
let incoming_tx_inner = self.incoming_tx.clone();
|
|
let connections_inner = self.connections.clone();
|
|
let conn_clone = conn.clone();
|
|
|
|
tokio::spawn(async move {
|
|
loop {
|
|
match conn_clone.read_datagram().await {
|
|
Ok(bytes) => {
|
|
if let Err(_) = incoming_tx_inner.send((peer_id, bytes)).await {
|
|
break;
|
|
}
|
|
}
|
|
Err(_) => {
|
|
connections_inner.lock().await.remove(&peer_id);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
});
|
|
|
|
conn
|
|
};
|
|
|
|
conn.send_datagram(data)
|
|
.map_err(|e| NetError::Connection(e.to_string()))?;
|
|
Ok(())
|
|
}
|
|
|
|
async fn receive_datagrams(&self) -> Result<Receiver<(EndpointId, Bytes)>, NetError> {
|
|
let mut rx_guard = self.incoming_rx.lock().await;
|
|
if let Some(rx) = rx_guard.take() {
|
|
Ok(rx)
|
|
} else {
|
|
Err(NetError::Other("Datagram receiver already subscribed".to_string()))
|
|
}
|
|
}
|
|
}
|