oxframe_preview/
server.rs

1//! Native preview server: WebSocket on `127.0.0.1`, protocol JSON for
2//! commands/events, binary frames (16-byte header + JPEG).
3//!
4//! The engine is single-threaded (`!Send`), so it lives on its own thread;
5//! the async side only moves messages. Bursts of `Seek` (scrubbing) are
6//! coalesced: only the latest seek in a batch is rendered.
7
8use std::io::Write;
9use std::sync::{Arc, Mutex, mpsc};
10use std::time::{Duration, Instant};
11
12use axum::Router;
13use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
14use axum::extract::{Query, State};
15use axum::http::{HeaderMap, StatusCode};
16use axum::response::{IntoResponse, Response};
17use axum::routing::get;
18use image::ExtendedColorType;
19use image::codecs::jpeg::JpegEncoder;
20use oxframe_core::Image;
21use oxframe_protocol::{Command, FrameFormat};
22use tokio::sync::mpsc::{UnboundedSender, unbounded_channel};
23
24use crate::{Encoder, Output, PreviewEngine, Session};
25
26/// Options of the native preview server.
27#[derive(Debug, Clone)]
28pub struct ServerOptions {
29    /// 0 picks a free port; the chosen one is printed as the first line of
30    /// stdout.
31    pub port: u16,
32    /// If set, clients must connect to `/ws?token=<token>`.
33    pub token: Option<String>,
34    /// Size limit of the encoded-frame cache, in MiB.
35    pub cache_mb: usize,
36    /// JPEG quality of preview frames, 1..=100.
37    pub jpeg_quality: u8,
38}
39
40impl Default for ServerOptions {
41    fn default() -> Self {
42        Self {
43            port: 0,
44            token: std::env::var("OXFRAME_PREVIEW_TOKEN")
45                .ok()
46                .filter(|t| !t.is_empty()),
47            cache_mb: 256,
48            jpeg_quality: 80,
49        }
50    }
51}
52
53/// Encodes preview frames as JPEG.
54pub fn jpeg_encoder(quality: u8) -> Encoder {
55    Box::new(move |img: &Image| {
56        let mut rgb = Vec::with_capacity((img.width * img.height * 3) as usize);
57        for px in img.data.chunks_exact(4) {
58            // Composite over black; preview frames are normally opaque.
59            let a = u16::from(px[3]);
60            rgb.extend(
61                px[..3]
62                    .iter()
63                    .map(|&c| ((u16::from(c) * a + 127) / 255) as u8),
64            );
65        }
66        let mut out = Vec::new();
67        JpegEncoder::new_with_quality(&mut out, quality)
68            .encode(&rgb, img.width, img.height, ExtendedColorType::Rgb8)
69            .expect("JPEG encoding into memory");
70        (FrameFormat::Jpeg, out)
71    })
72}
73
74enum Inbound {
75    Command(Command),
76}
77
78#[derive(Clone)]
79struct AppState {
80    commands: mpsc::Sender<Inbound>,
81    /// Output channel of the current client; a new connection replaces it.
82    client: Arc<Mutex<Option<UnboundedSender<Output>>>>,
83    token: Option<String>,
84}
85
86/// Runs the server until `Shutdown` or the process is killed. `make_engine`
87/// runs on the engine thread (engines are not `Send`).
88///
89/// # Errors
90///
91/// If the port cannot be bound.
92pub fn serve<E, F>(make_engine: F, opts: ServerOptions) -> std::io::Result<()>
93where
94    E: PreviewEngine + 'static,
95    F: FnOnce() -> Result<E, String> + Send + 'static,
96{
97    let rt = tokio::runtime::Builder::new_multi_thread()
98        .worker_threads(2)
99        .enable_all()
100        .build()?;
101    rt.block_on(async move {
102        let listener = tokio::net::TcpListener::bind(("127.0.0.1", opts.port)).await?;
103        let port = listener.local_addr()?.port();
104
105        let (cmd_tx, cmd_rx) = mpsc::channel::<Inbound>();
106        let client: Arc<Mutex<Option<UnboundedSender<Output>>>> = Arc::default();
107        let (ready_tx, ready_rx) = std::sync::mpsc::channel::<Result<(), String>>();
108        let engine_client = client.clone();
109        let (cache_mb, quality) = (opts.cache_mb, opts.jpeg_quality);
110        std::thread::Builder::new()
111            .name("oxframe-engine".into())
112            .spawn(move || {
113                let engine = match make_engine() {
114                    Ok(e) => e,
115                    Err(e) => {
116                        let _ = ready_tx.send(Err(e));
117                        return;
118                    }
119                };
120                let _ = ready_tx.send(Ok(()));
121                let session = Session::new(engine, jpeg_encoder(quality), cache_mb);
122                engine_loop(session, cmd_rx, engine_client);
123            })?;
124        // Fail before announcing the port if the project cannot load.
125        match ready_rx.recv() {
126            Ok(Ok(())) => {}
127            Ok(Err(e)) => return Err(std::io::Error::other(e)),
128            Err(_) => return Err(std::io::Error::other("engine thread died")),
129        }
130
131        // Contract with oxframe-cli: the first stdout line is the port.
132        let mut stdout = std::io::stdout();
133        writeln!(stdout, "{port}")?;
134        stdout.flush()?;
135
136        let state = AppState {
137            commands: cmd_tx,
138            client,
139            token: opts.token,
140        };
141        let app = Router::new()
142            .route("/ws", get(ws_handler))
143            .route("/healthz", get(|| async { "ok" }))
144            .with_state(state);
145        axum::serve(listener, app).await
146    })
147}
148
149fn engine_loop<E: PreviewEngine>(
150    mut session: Session<E>,
151    rx: mpsc::Receiver<Inbound>,
152    client: Arc<Mutex<Option<UnboundedSender<Output>>>>,
153) {
154    let clock = Instant::now();
155    let send = |outputs: Vec<Output>| {
156        if outputs.is_empty() {
157            return;
158        }
159        if let Some(tx) = client.lock().expect("client lock").as_ref() {
160            for o in outputs {
161                let _ = tx.send(o);
162            }
163        }
164    };
165    loop {
166        let now = clock.elapsed().as_secs_f64();
167        let wait = session
168            .next_tick_in(now)
169            .map_or(Duration::from_secs(3600), Duration::from_secs_f64);
170        let mut batch = Vec::new();
171        match rx.recv_timeout(wait) {
172            Ok(Inbound::Command(c)) => batch.push(c),
173            Err(mpsc::RecvTimeoutError::Timeout) => {}
174            Err(mpsc::RecvTimeoutError::Disconnected) => return,
175        }
176        while let Ok(Inbound::Command(c)) = rx.try_recv() {
177            batch.push(c);
178        }
179        let last_seek = batch
180            .iter()
181            .rposition(|c| matches!(c, Command::Seek { .. }));
182        for (i, cmd) in batch.into_iter().enumerate() {
183            if matches!(cmd, Command::Seek { .. }) && Some(i) != last_seek {
184                continue;
185            }
186            if matches!(cmd, Command::Shutdown) {
187                std::process::exit(0);
188            }
189            let now = clock.elapsed().as_secs_f64();
190            send(session.handle(cmd, now));
191        }
192        send(session.tick(clock.elapsed().as_secs_f64()));
193    }
194}
195
196#[derive(serde::Deserialize)]
197struct WsQuery {
198    token: Option<String>,
199}
200
201/// Browsers send `Origin`; only loopback pages may connect. Non-browser
202/// clients (the CLI proxy) send none.
203fn origin_allowed(headers: &HeaderMap) -> bool {
204    let Some(origin) = headers.get("origin").and_then(|o| o.to_str().ok()) else {
205        return true;
206    };
207    let host = origin
208        .strip_prefix("http://")
209        .unwrap_or("")
210        .split(':')
211        .next()
212        .unwrap_or("");
213    matches!(host, "127.0.0.1" | "localhost")
214}
215
216async fn ws_handler(
217    ws: WebSocketUpgrade,
218    headers: HeaderMap,
219    Query(q): Query<WsQuery>,
220    State(state): State<AppState>,
221) -> Response {
222    if !origin_allowed(&headers) {
223        return (StatusCode::FORBIDDEN, "origin not allowed").into_response();
224    }
225    if state.token.is_some() && q.token != state.token {
226        return (StatusCode::UNAUTHORIZED, "bad token").into_response();
227    }
228    ws.on_upgrade(move |socket| client(socket, state))
229}
230
231async fn client(mut socket: WebSocket, state: AppState) {
232    let (tx, mut rx) = unbounded_channel::<Output>();
233    *state.client.lock().expect("client lock") = Some(tx);
234    loop {
235        tokio::select! {
236            out = rx.recv() => {
237                let Some(out) = out else { break };
238                let msg = match out {
239                    Output::Event(e) => Message::Text(serde_json::to_string(&e).expect("event json").into()),
240                    Output::Frame(bytes) => Message::Binary(bytes.into()),
241                };
242                if socket.send(msg).await.is_err() {
243                    break;
244                }
245            }
246            msg = socket.recv() => {
247                match msg {
248                    Some(Ok(Message::Text(t))) => match serde_json::from_str::<Command>(&t) {
249                        Ok(cmd) => {
250                            if state.commands.send(Inbound::Command(cmd)).is_err() {
251                                break;
252                            }
253                        }
254                        Err(e) => {
255                            let err = oxframe_protocol::Event::Error {
256                                kind: oxframe_protocol::ErrorKind::Runtime,
257                                message: format!("bad command: {e}"),
258                                spans: Vec::new(),
259                            };
260                            let _ = socket.send(Message::Text(serde_json::to_string(&err).expect("json").into())).await;
261                        }
262                    },
263                    Some(Ok(Message::Close(_))) | None | Some(Err(_)) => break,
264                    Some(Ok(_)) => {}
265                }
266            }
267        }
268    }
269}
oxframe_preview/server.rs

oxframe_preview/
server.rs
