sim/src/main.rs

mod cluster;
mod storage;
mod world;

use axum::Router;
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::extract::{DefaultBodyLimit, Path, Request, State};
use axum::http::{StatusCode, header};
use axum::middleware::{self, Next};
use axum::response::{Html, IntoResponse, Response};
use axum::routing::{get, post};
use cluster::gateway::{GatewayOptions, Msg, Outputs};
use cluster::shard::ShardOptions;
use include_dir::{Dir, include_dir};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::sync::{mpsc, oneshot, watch};
use world::{Shard, SimScript};

// The UI bundle that build.rs produced from ui/.
static UI: Dir = include_dir!("$OUT_DIR/ui-pkg");

// Minimal shell; the whole UI is the Leptos app in ui/ compiled to wasm.
const INDEX: &str = r#"<!doctype html>
<html><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1"><title>World Simulation</title></head>
<body style="background:#111;color:#ddd;font:13px monospace"><script type="module">
try { const m = await import('/pkg/ui.js'); await m.default({ module_or_path: '/pkg/ui_bg.wasm' }); }
catch (e) { document.body.textContent = 'Failed to load the UI: ' + e; }
</script></body></html>"#;

#[derive(Clone)]
struct App {
    terrain: Arc<Mutex<String>>,
    script: Arc<Mutex<String>>,
    state: watch::Receiver<Arc<str>>,
    tx: mpsc::UnboundedSender<Msg>,
}

/// Rejects cross-site requests (including WebSocket upgrades), which could otherwise edit the world.
async fn same_origin(req: Request, next: Next) -> Response {
    let h = req.headers();
    if let Some(origin) = h.get(header::ORIGIN).and_then(|v| v.to_str().ok()) {
        let host = h
            .get(header::HOST)
            .and_then(|v| v.to_str().ok())
            .unwrap_or("");
        let authority = origin.split_once("://").map_or(origin, |(_, rest)| rest);
        if authority != host {
            return StatusCode::FORBIDDEN.into_response();
        }
    }
    next.run(req).await
}

fn json(body: String) -> impl IntoResponse {
    ([(header::CONTENT_TYPE, "application/json")], body)
}

async fn index() -> Html<&'static str> {
    Html(INDEX)
}

async fn pkg(Path(path): Path<String>) -> Response {
    let Some(file) = UI.get_file(&path) else {
        return StatusCode::NOT_FOUND.into_response();
    };
    let ctype = match path.rsplit('.').next() {
        Some("js") => "text/javascript",
        Some("wasm") => "application/wasm",
        _ => "application/octet-stream",
    };
    ([(header::CONTENT_TYPE, ctype)], file.contents()).into_response()
}

async fn terrain(State(app): State<App>) -> impl IntoResponse {
    json(app.terrain.lock().unwrap().clone())
}

async fn state(State(app): State<App>) -> impl IntoResponse {
    json(app.state.borrow().to_string())
}

async fn get_script(State(app): State<App>) -> String {
    app.script.lock().unwrap().clone()
}

async fn post_script(State(app): State<App>, body: String) -> Response {
    let (rtx, rrx) = oneshot::channel();
    let _ = app.tx.send(Msg::Script(body, rtx));
    match tokio::time::timeout(Duration::from_secs(5), rrx).await {
        Ok(Ok(Ok(()))) => "ok".into_response(),
        Ok(Ok(Err(e))) => (StatusCode::BAD_REQUEST, e).into_response(),
        _ => (StatusCode::SERVICE_UNAVAILABLE, "timed out").into_response(),
    }
}

async fn post_reset(State(app): State<App>) -> Response {
    restart(app, false).await
}

async fn post_regenerate(State(app): State<App>) -> Response {
    restart(app, true).await
}

async fn restart(app: App, regenerate: bool) -> Response {
    let (reply, result) = oneshot::channel();
    if app.tx.send(Msg::Restart { regenerate, reply }).is_err() {
        return StatusCode::SERVICE_UNAVAILABLE.into_response();
    }
    match tokio::time::timeout(Duration::from_secs(5), result).await {
        Ok(Ok(seed)) => seed.to_string().into_response(),
        _ => StatusCode::SERVICE_UNAVAILABLE.into_response(),
    }
}

async fn post_edit(State(app): State<App>, body: String) -> StatusCode {
    if body.len() > 256 {
        return StatusCode::PAYLOAD_TOO_LARGE;
    }
    let _ = app.tx.send(Msg::Edit(body));
    StatusCode::OK
}

async fn ws(upgrade: WebSocketUpgrade, State(app): State<App>) -> Response {
    upgrade.on_upgrade(move |socket| client(socket, app))
}

#[cfg(feature = "console")]
#[derive(serde::Deserialize)]
struct ConsoleTarget {
    #[serde(default)]
    shard: usize,
}

#[cfg(feature = "console")]
async fn picard_ws(
    upgrade: WebSocketUpgrade,
    axum::extract::Query(target): axum::extract::Query<ConsoleTarget>,
    State(app): State<App>,
) -> Response {
    upgrade.on_upgrade(move |socket| picard_client(socket, app, target.shard))
}

#[cfg(feature = "console")]
async fn picard_client(mut socket: WebSocket, app: App, shard: usize) {
    while let Some(message) = socket.recv().await {
        match message {
            Ok(Message::Text(text)) if text.len() <= picard::console::MAX_REQUEST_BYTES => {
                let (reply, result) = oneshot::channel();
                if app
                    .tx
                    .send(Msg::PicardConsole {
                        request: text.to_string(),
                        shard,
                        reply,
                    })
                    .is_err()
                {
                    return;
                }
                let response = match tokio::time::timeout(Duration::from_secs(5), result).await {
                    Ok(Ok(response)) => response,
                    Ok(Err(_)) | Err(_) => {
                        r#"{"id":null,"error":{"code":"unavailable","message":"Picard console did not respond"}}"#
                            .to_owned()
                    }
                };
                if socket.send(Message::Text(response.into())).await.is_err() {
                    return;
                }
            }
            Ok(Message::Text(_)) => {
                if socket
                    .send(Message::Text(
                        r#"{"id":null,"error":{"code":"invalid_request","message":"Request exceeds 65536 bytes"}}"#
                            .into(),
                    ))
                    .await
                    .is_err()
                {
                    return;
                }
            }
            Ok(Message::Close(_)) | Err(_) | Ok(Message::Binary(_)) => return,
            Ok(Message::Ping(_) | Message::Pong(_)) => {}
        }
    }
}

/// Pushes every new snapshot to the client; text frames from the client are edit commands.
async fn client(mut socket: WebSocket, app: App) {
    let mut rx = app.state.clone();
    loop {
        let snapshot = rx.borrow_and_update().to_string();
        if socket.send(Message::Text(snapshot.into())).await.is_err() {
            return;
        }
        loop {
            tokio::select! {
                changed = rx.changed() => {
                    if changed.is_err() {
                        return;
                    }
                    break;
                }
                msg = socket.recv() => match msg {
                    Some(Ok(Message::Text(t))) if t.len() <= 256 => {
                        let _ = app.tx.send(Msg::Edit(t.as_str().to_owned()));
                    }
                    Some(Ok(Message::Close(_)) | Err(_)) | None => return,
                    _ => {}
                },
            }
        }
    }
}

fn load_script() -> (String, SimScript) {
    let config_path = std::env::var("SIM_CONFIG").unwrap_or_else(|_| "config.rn".into());
    let config_source = std::fs::read_to_string(&config_path)
        .unwrap_or_else(|e| panic!("failed to read simulation config at {config_path}: {e}"));
    let sim_script = SimScript::load(&config_source)
        .unwrap_or_else(|e| panic!("failed to load simulation config at {config_path}: {e}"));
    (config_source, sim_script)
}

/// Parses `shard --index I --count N --seed S --database URL [--listen ADDR] [--attached]`.
fn shard_options(args: &[String], script: SimScript) -> ShardOptions {
    let mut index = None;
    let mut count = None;
    let mut seed = 42;
    let mut cluster_token = None;
    let mut database_url = None;
    let mut listen = "127.0.0.1:0".to_owned();
    let mut attached = false;
    let mut args = args.iter();
    while let Some(flag) = args.next() {
        if flag == "--attached" {
            attached = true;
            continue;
        }
        let value = args
            .next()
            .unwrap_or_else(|| panic!("missing value for {flag}"));
        let number = || {
            value
                .parse()
                .unwrap_or_else(|_| panic!("invalid value for {flag}: {value}"))
        };
        match flag.as_str() {
            "--index" => index = Some(number()),
            "--count" => count = Some(number()),
            "--seed" => {
                seed = value
                    .parse()
                    .unwrap_or_else(|_| panic!("invalid seed: {value}"))
            }
            "--cluster-token" => cluster_token = Some(value.clone()),
            "--database" => database_url = Some(value.clone()),
            "--listen" => listen.clone_from(value),
            _ => panic!("unknown shard option {flag}"),
        }
    }
    let shard = Shard::new(
        index.expect("--index is required"),
        count.expect("--count is required"),
    )
    .expect("--index must be less than --count");
    ShardOptions {
        shard,
        seed,
        cluster_token: cluster_token.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
        database_url: database_url.expect("--database is required"),
        listen,
        script,
        attached,
    }
}

#[tokio::main]
async fn main() {
    let args: Vec<String> = std::env::args().skip(1).collect();
    let is_shard = args.first().is_some_and(|mode| mode == "shard");
    let filter = tracing_subscriber::EnvFilter::try_from_default_env()
        .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("picard=info,sim=info"));
    // Shards keep stdout for announcing their address to the gateway.
    tracing_subscriber::fmt()
        .with_env_filter(filter)
        .with_target(false)
        .with_writer(std::io::stderr)
        .init();

    let (config_source, sim_script) = load_script();
    if is_shard {
        cluster::shard::run(shard_options(&args[1..], sim_script)).await;
        return;
    }

    let seed = args.first().and_then(|s| s.parse().ok()).unwrap_or(42u64);
    let shards = std::env::var("SIM_SHARDS").ok().map_or(3, |count| {
        count
            .parse::<usize>()
            .ok()
            .filter(|&count| count > 0)
            .unwrap_or_else(|| panic!("SIM_SHARDS must be a positive number, not {count}"))
    });
    let database_template = std::env::var("SHARD_DATABASE_URL")
        .unwrap_or_else(|_| "sqlite://sim-shard-{shard}.db?mode=rwc".into());
    assert!(
        shards == 1 || database_template.contains("{shard}"),
        "SHARD_DATABASE_URL must contain {{shard}} so that every shard gets its own database"
    );
    let addr = std::env::var("SIM_ADDR").unwrap_or_else(|_| "127.0.0.1:7878".into());
    let listener = tokio::net::TcpListener::bind(&addr)
        .await
        .expect("failed to bind");

    let terrain_buf = Arc::new(Mutex::new(String::new()));
    let script = Arc::new(Mutex::new(config_source));
    let (tx, rx) = mpsc::unbounded_channel();
    let (state_tx, state_rx) = watch::channel(Arc::<str>::from("{}"));
    {
        let outputs = Outputs {
            terrain: Arc::clone(&terrain_buf),
            script: Arc::clone(&script),
            state: state_tx,
        };
        let options = GatewayOptions {
            shards,
            seed,
            database_template,
        };
        std::thread::Builder::new()
            .name("gateway".into())
            .spawn(move || cluster::gateway::conduct(&options, &outputs, rx))
            .expect("failed to start the gateway");
    }

    let app = App {
        terrain: terrain_buf,
        script,
        state: state_rx,
        tx,
    };
    let router = Router::new()
        .route("/", get(index))
        .route("/api/terrain", get(terrain))
        .route("/api/state", get(state))
        .route("/api/script", get(get_script).post(post_script))
        .route("/api/reset", post(post_reset))
        .route("/api/regenerate", post(post_regenerate))
        .route("/api/edit", post(post_edit))
        .route("/ws", get(ws))
        .route("/pkg/{*path}", get(pkg));
    #[cfg(feature = "console")]
    let router = router.route("/api/actors/ws", get(picard_ws));
    let router = router
        .layer(DefaultBodyLimit::max(65536))
        .layer(middleware::from_fn(same_origin))
        .with_state(app);

    println!("Simulation running at http://{addr} with {shards} shards");
    axum::serve(listener, router).await.unwrap();
}