sim/src/main.rs

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 include_dir::{Dir, include_dir};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use storage::Storage;
use tokio::sync::{mpsc, oneshot, watch};
use world::{DEFAULT_SCRIPT, World};

// 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>"#;

enum Msg {
    Edit(String),
    Script(String, oneshot::Sender<Result<(), String>>),
    Restart {
        regenerate: bool,
        reply: oneshot::Sender<u64>,
    },
}

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

/// The Rune VM is not Send, so the world lives on this thread and publishes JSON snapshots.
fn run_sim(
    seed: u64,
    snapshot: Option<String>,
    mut rx: mpsc::UnboundedReceiver<Msg>,
    terrain: Arc<Mutex<String>>,
    script: Arc<Mutex<String>>,
    state_tx: watch::Sender<Arc<str>>,
    snapshot_tx: watch::Sender<Option<String>>,
) {
    let mut world = match snapshot {
        Some(snapshot) => World::from_snapshot(&snapshot, seed)
            .unwrap_or_else(|e| panic!("failed to restore simulation snapshot: {e}")),
        None => World::new(seed),
    };
    *script.lock().unwrap() = world.script().to_string();
    let mut rev = u32::MAX;
    let mut last_saved = Instant::now();
    snapshot_tx.send_replace(Some(
        world
            .snapshot_json()
            .expect("failed to serialize initial simulation snapshot"),
    ));
    loop {
        while let Ok(msg) = rx.try_recv() {
            match msg {
                Msg::Edit(cmd) => world.apply(&cmd),
                Msg::Script(src, reply) => {
                    let r = world.set_script(&src);
                    if r.is_ok() {
                        *script.lock().unwrap() = src;
                    }
                    let _ = reply.send(r);
                }
                Msg::Restart { regenerate, reply } => {
                    let seed = if regenerate {
                        world.regenerated_seed()
                    } else {
                        world.seed()
                    };
                    world = world.restarted(seed);
                    *script.lock().unwrap() = world.script().to_string();
                    rev = u32::MAX;
                    snapshot_tx.send_replace(Some(
                        world
                            .snapshot_json()
                            .expect("failed to serialize restarted simulation"),
                    ));
                    let _ = reply.send(seed);
                }
            }
        }
        if !world.paused {
            for _ in 0..world.speed {
                world.step();
            }
        }
        if world.terrain_rev != rev {
            rev = world.terrain_rev;
            *terrain.lock().unwrap() = world.terrain_json();
        }
        state_tx.send_replace(Arc::from(world.state_json()));
        if last_saved.elapsed() >= Duration::from_secs(1) {
            snapshot_tx.send_replace(Some(
                world
                    .snapshot_json()
                    .expect("failed to serialize simulation snapshot"),
            ));
            last_saved = Instant::now();
        }
        std::thread::sleep(Duration::from_millis(100));
    }
}

/// 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))
}

/// 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(_))) | Some(Err(_)) | None => return,
                    _ => {}
                },
            }
        }
    }
}

#[tokio::main]
async fn main() {
    let seed = std::env::args()
        .nth(1)
        .and_then(|s| s.parse().ok())
        .unwrap_or(42u64);
    let addr = std::env::var("SIM_ADDR").unwrap_or_else(|_| "127.0.0.1:7878".into());
    let database_url =
        std::env::var("DATABASE_URL").unwrap_or_else(|_| "sqlite://sim.db?mode=rwc".into());
    let storage = Storage::connect(&database_url)
        .await
        .expect("failed to connect to simulation database");
    let saved_snapshot = storage
        .load()
        .await
        .expect("failed to load simulation snapshot");
    let (snapshot_tx, snapshot_rx) = watch::channel(None::<String>);
    let snapshot_storage = storage.clone();
    tokio::spawn(async move {
        let mut snapshots = snapshot_rx;
        let mut interval = tokio::time::interval(Duration::from_secs(1));
        interval.tick().await;
        loop {
            tokio::select! {
                _ = interval.tick() => {}
                changed = snapshots.changed() => {
                    if changed.is_err() {
                        break;
                    }
                }
            }
            let Some(snapshot) = snapshots.borrow_and_update().clone() else {
                continue;
            };
            if let Err(error) = snapshot_storage.save(&snapshot).await {
                eprintln!("failed to persist simulation snapshot: {error}");
            }
        }
    });

    let terrain_buf = Arc::new(Mutex::new(String::new()));
    let script = Arc::new(Mutex::new(DEFAULT_SCRIPT.to_string()));
    let (tx, rx) = mpsc::unbounded_channel();
    let (state_tx, state_rx) = watch::channel(Arc::<str>::from("{}"));
    {
        let (terrain_buf, script) = (Arc::clone(&terrain_buf), Arc::clone(&script));
        std::thread::spawn(move || {
            run_sim(
                seed,
                saved_snapshot,
                rx,
                terrain_buf,
                script,
                state_tx,
                snapshot_tx,
            )
        });
    }

    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))
        .layer(DefaultBodyLimit::max(65536))
        .layer(middleware::from_fn(same_origin))
        .with_state(app);

    let listener = tokio::net::TcpListener::bind(&addr)
        .await
        .expect("failed to bind");
    println!("Simulation running at http://{addr}");
    axum::serve(listener, router).await.unwrap();
}