worldwideportal-server/src/main.rs

use actix_web::{
    self,
    error::InternalError,
    get,
    http::StatusCode,
    middleware::Logger,
    rt::{self, net::TcpStream},
    web::{self, Data},
    App, Error, HttpRequest, HttpResponse, HttpServer, Responder,
};
use actix_ws::{AggregatedMessage, CloseCode, Closed};
use anyhow::Context;
use futures_util::StreamExt;
use serde_json::json;
use std::{env, fs, num::ParseIntError};
use tokio::io::{AsyncReadExt, AsyncWriteExt};

pub struct ServerConfig {
    pub upstream_mud: String,
    pub banner_to_mud: String,
    pub listen_port: u16,
    pub bind_address: String,
    pub startup_script_file: String,
}

#[get("/ws")]
async fn ws(
    config_data: Data<ServerConfig>,
    req: HttpRequest,
    body: web::Payload,
) -> impl Responder {
    let (response, mut session, stream) = actix_ws::handle(&req, body)?;
    let mut stream = stream.aggregate_continuations().max_continuation_size(1024);

    let mut tcp_stream: TcpStream = TcpStream::connect(&config_data.upstream_mud).await?;
    let subst_banner = config_data
        .banner_to_mud
        .replace(
            "%i",
            &req.peer_addr()
                .map(|a| a.to_string())
                .unwrap_or_else(|| "unknown".to_owned()),
        )
        .replace("%n", "\r\n");
    tcp_stream.write_all(subst_banner.as_bytes()).await?;

    let script = fetch_startup_script(&config_data.startup_script_file)
        .map_err(|e| InternalError::new(e, StatusCode::INTERNAL_SERVER_ERROR))?;

    rt::spawn(async move {
        if session
            .text(
                json!({
                    "RunLua": script
                })
                .to_string(),
            )
            .await
            .is_err()
        {
            let _ = session.close(Some(CloseCode::Normal.into())).await;
            return;
        }
        let mut readbuf: [u8; 1024] = [0; 1024];
        loop {
            tokio::select! {
                ws_msg = stream.next() => {
                    match ws_msg {
                        None => break,
                        Some(Err(_e)) => break,
                        Some(Ok(AggregatedMessage::Binary(bin))) => {
                            if tcp_stream.write_all(&bin).await.is_err() {
                                break
                            }
                        }

                        Some(Ok(AggregatedMessage::Ping(msg))) => {
                            if let Err(Closed) = session.pong(&msg).await {
                                break
                            }
                        }
                        Some(Ok(_)) => {},
                    }
                },
                tcp_data_len = tcp_stream.read(&mut readbuf) => {
                    match tcp_data_len {
                        Err(_e) => break,
                        Ok(0) => break,
                        Ok(n) =>
                            if let Err(Closed) = session.binary(readbuf[0..n].to_vec()).await {
                                break;
                            }
                    }

                }
            }
        }
        let _ = session.close(Some(CloseCode::Normal.into())).await;
    });

    Ok::<HttpResponse, Error>(response)
}

fn extract_server_config_from_environment() -> anyhow::Result<ServerConfig> {
    Ok(ServerConfig {
        upstream_mud: env::var("UPSTREAM_MUD").map_err(|_| {
            anyhow::Error::msg(
                "Expected UPSTREAM_MUD environment variable specifying where to connect.",
            )
        })?,
        banner_to_mud: env::var("BANNER_TO_MUD").map_err(|_| {
            anyhow::Error::msg(
                "Expected BANNER_TO_MUD environment variable specifying message to send to MUD.",
            )
        })?,
        listen_port: env::var("LISTEN_PORT")
            .map_err(|_| {
                anyhow::Error::msg(
                    "Expected LISTEN_PORT environment variable specifying port to listen on.",
                )
            })
            .and_then(|v| {
                v.parse::<u16>().map_err(|_e: ParseIntError| {
                    anyhow::Error::msg("LISTEN_PORT should be a decimal port number")
                })
            })?,
        bind_address: env::var("BIND_ADDRESS").unwrap_or_else(|_| "::".to_owned()),
        startup_script_file: env::var("STARTUP_SCRIPT_FILE")
            .map_err(|_| {
                anyhow::Error::msg(
                    "Expected STARTUP_SCRIPT_FILE environment variable containing filename of script to send to client.",
                )
            })?
,
    })
}

// We load this on each connection so it can change.
fn fetch_startup_script(filename: &str) -> anyhow::Result<String> {
    Ok(fs::read_to_string(filename)?)
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    env_logger::init();

    let config = extract_server_config_from_environment()?;

    let data = Data::new(config);
    let config = data.get_ref();

    // Do this early so we fail fast if config is wrong.
    fetch_startup_script(&config.startup_script_file)
        .context("While checking STARTUP_SCRIPT_FILE can be read")?;

    let server_data = data.clone();
    HttpServer::new(move || {
        let logger = Logger::default();
        App::new()
            .wrap(logger)
            .app_data(server_data.clone())
            .service(ws)
    })
    .bind((config.bind_address.clone(), config.listen_port))?
    .run()
    .await?;
    Ok(())
}