use std::net::SocketAddr; use prometheus::{IntCounter, IntGauge, Registry as MetricsRegistry}; use serde::Deserialize; use tokio::io::{AsyncBufReadExt, AsyncWrite, AsyncWriteExt, BufReader, BufWriter}; use tokio::net::tcp::{ReadHalf, WriteHalf}; use tokio::net::{TcpListener, TcpStream}; use tokio::sync::oneshot::channel; use crate::core::player::{ConnectionId, PlayerHandle, PlayerId, PlayerRegistry, Updates}; use crate::core::room::RoomId; use crate::prelude::*; use crate::protos::irc::client::{client_message, ClientMessage}; use crate::protos::irc::server::{ServerMessage, ServerMessageBody}; use crate::protos::irc::{Chan, Recipient}; use crate::util::Terminator; #[derive(Deserialize, Debug, Clone)] pub struct ServerConfig { pub listen_on: SocketAddr, pub server_name: String, } #[derive(Debug)] struct RegisteredUser { nickname: Vec, username: Vec, realname: Vec, } async fn handle_socket( config: ServerConfig, mut stream: TcpStream, socket_addr: SocketAddr, mut players: PlayerRegistry, ) -> Result<()> { let (reader, writer) = stream.split(); let mut reader: BufReader = BufReader::new(reader); let mut writer = BufWriter::new(writer); ServerMessage { tags: vec![], sender: Some(config.server_name.as_bytes().to_vec()), body: ServerMessageBody::Notice { first_target: b"*".to_vec(), rest_targets: vec![], text: b"Welcome to my server!".to_vec(), }, } .write_async(&mut writer) .await?; writer.flush().await?; let registered_user: Result = handle_registration(&mut reader, &mut writer).await; match registered_user { Ok(user) => { handle_registered_socket(config, socket_addr, players, &mut reader, &mut writer, user) .await?; } Err(_) => {} } stream.shutdown().await?; Ok(()) } async fn handle_registration<'a>( reader: &mut BufReader>, writer: &mut BufWriter>, ) -> Result { let mut buffer = vec![]; let mut future_nickname: Option> = None; let mut future_username: Option<(Vec, Vec)> = None; loop { let res = reader.read_until(b'\n', &mut buffer).await; match res { Ok(len) => { if len == 0 { log::info!("Terminating socket"); break Err(anyhow::Error::msg("EOF")); } } Err(err) => { log::warn!("Failed to read from socket: {err}"); break Err(err.into()); } } let parsed = client_message(&buffer[..]); match parsed { Ok((rest, msg)) => { log::info!("Incoming IRC message: {msg:?}"); match msg { ClientMessage::Nick { nickname } => { if let Some((username, realname)) = future_username { break Ok(RegisteredUser { nickname, username, realname, }); } else { future_nickname = Some(nickname); } } ClientMessage::User { username, realname } => { if let Some(nickname) = future_nickname { break Ok(RegisteredUser { nickname, username, realname, }); } else { future_username = Some((username, realname)); } } _ => {} } } Err(err) => { log::warn!("Failed to parse IRC message: {err}"); } } buffer.clear(); } } async fn handle_registered_socket<'a>( config: ServerConfig, socket_addr: SocketAddr, mut players: PlayerRegistry, reader: &mut BufReader>, writer: &mut BufWriter>, user: RegisteredUser, ) -> Result<()> { let mut buffer = vec![]; log::info!("Handling registered user: {user:?}"); let mut user_handle = players .get_or_create_player(PlayerId(user.nickname.clone())) .await; let (my_connection_id, mut connnection) = user_handle.subscribe().await; ServerMessage { tags: vec![], sender: Some(config.server_name.as_bytes().to_vec()), body: ServerMessageBody::N001Welcome { client: user.nickname.clone(), text: b"Welcome to Kek Server".to_vec(), }, } .write_async(writer) .await?; ServerMessage { tags: vec![], sender: Some(config.server_name.as_bytes().to_vec()), body: ServerMessageBody::N002YourHost { client: user.nickname.clone(), text: b"Welcome to Kek Server".to_vec(), }, } .write_async(writer) .await?; ServerMessage { tags: vec![], sender: Some(config.server_name.as_bytes().to_vec()), body: ServerMessageBody::N003Created { client: user.nickname.clone(), text: b"Welcome to Kek Server".to_vec(), }, } .write_async(writer) .await?; ServerMessage { tags: vec![], sender: Some(config.server_name.as_bytes().to_vec()), body: ServerMessageBody::N004MyInfo { client: user.nickname.clone(), hostname: config.server_name.as_bytes().to_vec(), softname: b"kek-0.1.alpha.3".to_vec(), }, } .write_async(writer) .await?; ServerMessage { tags: vec![], sender: Some(config.server_name.as_bytes().to_vec()), body: ServerMessageBody::N005ISupport { client: user.nickname.clone(), params: b"CHANTYPES=#".to_vec(), }, } .write_async(writer) .await?; writer.flush().await?; loop { select! { biased; len = reader.read_until(b'\n', &mut buffer) => { let len = len?; let len = if len == 0 { log::info!("EOF, Terminating socket"); break; } else { len }; handle_incoming_message(&buffer[0..len], &config, &user, &mut user_handle, &my_connection_id, writer).await?; buffer.clear(); }, update = connnection.recv() => { match update.unwrap() { Updates::RoomJoined { room_id } => {}, Updates::NewMessage { author_id, connection_id, room_id, body } => { if my_connection_id != connection_id { ServerMessage { tags: vec![], sender: Some(author_id.0.clone()), body: ServerMessageBody::PrivateMessage { target: Recipient::Chan(Chan::Global(room_id.0)), body: body.as_bytes().to_vec() } }.write_async(writer).await?; writer.flush().await? } }, } } } } Ok(()) } async fn handle_incoming_message( buffer: &[u8], config: &ServerConfig, user: &RegisteredUser, user_handle: &mut PlayerHandle, my_connection_id: &ConnectionId, writer: &mut (impl AsyncWrite + Unpin), ) -> Result<()> { let parsed = client_message(buffer); match parsed { Ok((rest, msg)) => match msg { ClientMessage::Ping { token } => { ServerMessage { tags: vec![], sender: None, body: ServerMessageBody::Pong { from: config.server_name.as_bytes().to_vec(), token, }, } .write_async(writer) .await?; writer.flush().await?; } ClientMessage::Join(ref chan) => { handle_join(&config, &user, user_handle, chan, writer).await?; } ClientMessage::PrivateMessage { recipient, body } => match recipient { Recipient::Chan(Chan::Global(room)) => match String::from_utf8(body) { Ok(body) => { user_handle .send_message( RoomId(room.clone()), my_connection_id.clone(), body.clone(), ) .await } Err(err) => log::warn!("failed to parse incoming message: {err}"), }, _ => log::warn!("Unsupported target type"), }, _ => {} }, Err(err) => { log::warn!("Failed to parse IRC message: {err}"); } } Ok(()) } async fn handle_join( config: &ServerConfig, user: &RegisteredUser, user_handle: &mut PlayerHandle, chan: &Chan, writer: &mut (impl AsyncWrite + Unpin), ) -> Result<()> { match chan { Chan::Global(ref room) => { user_handle.join_room(RoomId(room.clone())).await; ServerMessage { tags: vec![], sender: Some(user.nickname.clone()), body: ServerMessageBody::Join(chan.clone()), } .write_async(writer) .await?; ServerMessage { tags: vec![], sender: Some(config.server_name.as_bytes().to_vec()), body: ServerMessageBody::N332Topic { client: user.nickname.clone(), chat: chan.clone(), topic: b"chan topic lol".to_vec(), }, } .write_async(writer) .await?; ServerMessage { tags: vec![], sender: Some(config.server_name.as_bytes().to_vec()), body: ServerMessageBody::N353NamesReply { client: user.nickname.clone(), chan: chan.clone(), members: user.nickname.clone(), }, } .write_async(writer) .await?; ServerMessage { tags: vec![], sender: Some(config.server_name.as_bytes().to_vec()), body: ServerMessageBody::N366NamesReplyEnd { client: user.nickname.clone(), chan: chan.clone(), }, } .write_async(writer) .await?; writer.flush().await?; } Chan::Local(_) => {} }; Ok(()) } pub async fn launch( config: ServerConfig, players: PlayerRegistry, metrics: MetricsRegistry, ) -> Result { log::info!("Starting IRC projection"); let (signal, mut rx) = channel(); let current_connections = IntGauge::new("irc_current_connections", "Open and alive TCP connections")?; let total_connections = IntCounter::new( "irc_total_connections", "Total number of opened connections", )?; metrics.register(Box::new(current_connections.clone()))?; metrics.register(Box::new(total_connections.clone()))?; let listener = TcpListener::bind(config.listen_on).await?; log::debug!("Listener started"); let handle = tokio::task::spawn(async move { loop { select! { biased; _ = &mut rx => break, new_conn = listener.accept() => { match new_conn { Ok((stream, socket_addr)) => { let config = config.clone(); total_connections.inc(); current_connections.inc(); log::debug!("Incoming connection from {socket_addr}"); let players_clone = players.clone(); let current_connections_clone = current_connections.clone(); let handle = tokio::task::spawn(async move { match handle_socket(config, stream, socket_addr, players_clone).await { Ok(_) => log::info!("Connection terminated"), Err(err) => log::warn!("Connection failed: {err}"), } current_connections_clone.dec(); }); }, Err(err) => log::warn!("Failed to accept new connection: {err}"), } }, } } log::info!("Stopping IRC projection"); Ok(()) }); let terminator = Terminator::from_raw(signal, handle); Ok(terminator) }