103 lines
2.6 KiB
Rust
103 lines
2.6 KiB
Rust
use async_std::sync::Arc;
|
|
use async_std::sync::Mutex;
|
|
use std::collections::HashMap;
|
|
|
|
use async_std::net::{SocketAddr, TcpListener, TcpStream};
|
|
use async_std::sync::{Receiver, Sender};
|
|
use async_std::task;
|
|
use async_tungstenite::{accept_async, WebSocketStream};
|
|
use futures_util::{stream::SplitSink, stream::SplitStream, SinkExt, StreamExt};
|
|
|
|
#[derive(Debug)]
|
|
pub struct NetworkMessage {
|
|
pub conn_id: usize,
|
|
pub data: tungstenite::Message,
|
|
}
|
|
|
|
struct Connection {
|
|
stream: SplitSink<WebSocketStream<TcpStream>, tungstenite::Message>,
|
|
conn_id: usize,
|
|
}
|
|
|
|
pub struct NetworkManager {
|
|
connections: HashMap<usize, Connection>,
|
|
}
|
|
|
|
impl NetworkManager {
|
|
pub fn new() -> Self {
|
|
NetworkManager {
|
|
connections: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
fn register(&mut self, conn: Connection) {
|
|
self.connections.insert(conn.conn_id, conn);
|
|
}
|
|
|
|
async fn send(&self, message: NetworkMessage) {
|
|
let conn = self
|
|
.connections
|
|
.get(&message.conn_id)
|
|
.expect("cant send message to an unregistered connection");
|
|
conn.stream.send(message.data).await;
|
|
}
|
|
}
|
|
|
|
pub async fn listen(
|
|
net: Arc<NetworkManager>,
|
|
bind: String,
|
|
incoming: Sender<NetworkMessage>,
|
|
outgoing: Receiver<NetworkMessage>,
|
|
) {
|
|
let network = Arc::new(Mutex::new(NetworkManager::new()));
|
|
let listener = TcpListener::bind(&bind).await.expect("Can't listen");
|
|
log::info!("Listening on: {}", bind);
|
|
task::spawn(write_loop(network.clone(), outgoing));
|
|
while let Ok((stream, _)) = listener.accept().await {
|
|
let peer = stream
|
|
.peer_addr()
|
|
.expect("connected streams should have a peer address");
|
|
log::info!("Peer address: {}", peer);
|
|
let conn = accept_connection(peer, stream, incoming.clone()).await;
|
|
network.lock().await.register(conn);
|
|
}
|
|
}
|
|
|
|
async fn accept_connection(
|
|
peer: SocketAddr,
|
|
stream: TcpStream,
|
|
incoming: Sender<NetworkMessage>,
|
|
) -> Connection {
|
|
let conn_id = rand::random();
|
|
let ws_stream = accept_async(stream).await.expect("could not accept");
|
|
log::info!("New WebSocket connection: {}", peer);
|
|
let (ws_sender, ws_receiver) = ws_stream.split();
|
|
task::spawn(read_loop(ws_receiver, incoming, conn_id));
|
|
Connection {
|
|
stream: ws_sender,
|
|
conn_id,
|
|
}
|
|
}
|
|
|
|
async fn read_loop(
|
|
mut stream: SplitStream<WebSocketStream<TcpStream>>,
|
|
ch: Sender<NetworkMessage>,
|
|
conn_id: usize,
|
|
) {
|
|
loop {
|
|
let data = stream
|
|
.next()
|
|
.await
|
|
.expect("failed getting the next message")
|
|
.expect("received error while reading");
|
|
ch.send(NetworkMessage { conn_id, data }).await;
|
|
}
|
|
}
|
|
|
|
async fn write_loop(net: Arc<Mutex<NetworkManager>>, ch: Receiver<NetworkMessage>) {
|
|
loop {
|
|
let msg = ch.recv().await.expect("failed to receive outgoing message");
|
|
net.lock().await.send(msg).await;
|
|
}
|
|
}
|