ps_blitz_debug_control/
lib.rs

1//! Debug-only, loopback WebDriver transport for Blitz renderers.
2//!
3//! The server owns networking and session authentication. Renderer commands
4//! cross a serialized channel and must be executed by the UI/runtime thread.
5
6use std::fs::{self, OpenOptions};
7use std::io::{self, Read, Write};
8use std::net::{SocketAddr, TcpListener, TcpStream};
9use std::path::PathBuf;
10use std::sync::Arc;
11use std::sync::atomic::{AtomicBool, Ordering};
12use std::sync::mpsc::{self, Receiver, RecvTimeoutError, SyncSender, TrySendError};
13use std::thread::{self, JoinHandle};
14use std::time::Duration;
15
16use serde::Serialize;
17use serde_json::{Value, json};
18
19const PROTOCOL_VERSION: u32 = 1;
20const MAX_REQUEST_BYTES: usize = 1024 * 1024;
21const COMMAND_TIMEOUT: Duration = Duration::from_secs(5);
22
23/// Configuration for a loopback debug-control server.
24#[derive(Debug, Clone)]
25pub struct ServerConfig {
26    /// Loopback address to bind. Port zero asks the OS to choose a free port.
27    pub bind_address: SocketAddr,
28    /// Atomically written once the server is accepting connections.
29    pub descriptor_path: PathBuf,
30    /// Git revision or build identifier for the renderer.
31    pub renderer_revision: String,
32}
33
34/// A command forwarded from the HTTP server to the renderer thread.
35#[derive(Debug)]
36pub struct ControlRequest {
37    pub method: String,
38    pub path: String,
39    pub body: Value,
40    reply: SyncSender<ControlResponse>,
41}
42
43impl ControlRequest {
44    /// Complete this request. Failure means the client already disconnected.
45    pub fn respond(self, response: ControlResponse) -> Result<(), ControlResponse> {
46        self.reply.send(response).map_err(|error| error.0)
47    }
48}
49
50/// A renderer response represented using W3C WebDriver success/error values.
51#[derive(Debug)]
52pub enum ControlResponse {
53    Success(Value),
54    Error {
55        error: String,
56        message: String,
57        stacktrace: String,
58    },
59}
60
61impl ControlResponse {
62    pub fn unsupported(message: impl Into<String>) -> Self {
63        Self::Error {
64            error: "unsupported operation".into(),
65            message: message.into(),
66            stacktrace: String::new(),
67        }
68    }
69}
70
71#[derive(Debug, Serialize)]
72#[serde(rename_all = "camelCase")]
73struct Descriptor<'a> {
74    pid: u32,
75    address: String,
76    token: &'a str,
77    protocol_version: u32,
78    renderer: &'static str,
79    renderer_revision: &'a str,
80}
81
82/// Running server. Dropping it shuts down the listener and removes discovery.
83pub struct DebugServer {
84    address: SocketAddr,
85    token: String,
86    descriptor_path: PathBuf,
87    shutdown: Arc<AtomicBool>,
88    thread: Option<JoinHandle<()>>,
89}
90
91impl DebugServer {
92    /// Bind a loopback port, write the descriptor, and start the server thread.
93    pub fn start(config: ServerConfig) -> io::Result<(Self, Receiver<ControlRequest>)> {
94        if !config.bind_address.ip().is_loopback() {
95            return Err(io::Error::new(
96                io::ErrorKind::InvalidInput,
97                "debug control must bind to a loopback address",
98            ));
99        }
100        let listener = TcpListener::bind(config.bind_address)?;
101        let address = listener.local_addr()?;
102        let token = random_hex(32)?;
103        let (command_tx, command_rx) = mpsc::sync_channel(1);
104        let shutdown = Arc::new(AtomicBool::new(false));
105
106        let thread_shutdown = Arc::clone(&shutdown);
107        let thread_token = token.clone();
108        let thread = thread::Builder::new()
109            .name("blitz-debug-control".into())
110            .spawn(move || server_loop(listener, &thread_token, command_tx, thread_shutdown))?;
111
112        if let Err(error) = write_descriptor(&config, address, &token) {
113            shutdown.store(true, Ordering::Release);
114            let _ = TcpStream::connect(address);
115            let _ = thread.join();
116            return Err(error);
117        }
118
119        Ok((
120            Self {
121                address,
122                token,
123                descriptor_path: config.descriptor_path,
124                shutdown,
125                thread: Some(thread),
126            },
127            command_rx,
128        ))
129    }
130
131    pub fn address(&self) -> SocketAddr {
132        self.address
133    }
134
135    pub fn token(&self) -> &str {
136        &self.token
137    }
138
139    pub fn shutdown(mut self) {
140        self.stop();
141    }
142
143    fn stop(&mut self) {
144        self.shutdown.store(true, Ordering::Release);
145        let _ = TcpStream::connect(self.address);
146        if let Some(thread) = self.thread.take() {
147            let _ = thread.join();
148        }
149        let _ = fs::remove_file(&self.descriptor_path);
150    }
151}
152
153impl Drop for DebugServer {
154    fn drop(&mut self) {
155        self.stop();
156    }
157}
158
159fn random_hex(byte_len: usize) -> io::Result<String> {
160    let mut bytes = vec![0; byte_len];
161    getrandom::fill(&mut bytes).map_err(io::Error::other)?;
162    let mut output = String::with_capacity(byte_len * 2);
163    for byte in bytes {
164        use std::fmt::Write as _;
165        write!(output, "{byte:02x}").unwrap();
166    }
167    Ok(output)
168}
169
170fn write_descriptor(config: &ServerConfig, address: SocketAddr, token: &str) -> io::Result<()> {
171    let descriptor = Descriptor {
172        pid: std::process::id(),
173        address: address.to_string(),
174        token,
175        protocol_version: PROTOCOL_VERSION,
176        renderer: "blitz",
177        renderer_revision: &config.renderer_revision,
178    };
179    let bytes = serde_json::to_vec_pretty(&descriptor).map_err(io::Error::other)?;
180    if let Some(parent) = config.descriptor_path.parent() {
181        fs::create_dir_all(parent)?;
182    }
183    let temporary = config
184        .descriptor_path
185        .with_extension(format!("tmp-{}", random_hex(8)?));
186    let result = (|| {
187        let mut options = OpenOptions::new();
188        options.write(true).create_new(true);
189        #[cfg(unix)]
190        {
191            use std::os::unix::fs::OpenOptionsExt;
192            options.mode(0o600);
193        }
194        let mut file = options.open(&temporary)?;
195        file.write_all(&bytes)?;
196        file.sync_all()?;
197        fs::rename(&temporary, &config.descriptor_path)
198    })();
199    if result.is_err() {
200        let _ = fs::remove_file(&temporary);
201    }
202    result
203}
204
205fn server_loop(
206    listener: TcpListener,
207    token: &str,
208    command_tx: SyncSender<ControlRequest>,
209    shutdown: Arc<AtomicBool>,
210) {
211    let mut active_session: Option<String> = None;
212    for connection in listener.incoming() {
213        if shutdown.load(Ordering::Acquire) {
214            break;
215        }
216        match connection {
217            Ok(mut stream) => {
218                let _ = stream.set_read_timeout(Some(COMMAND_TIMEOUT));
219                let _ = stream.set_write_timeout(Some(COMMAND_TIMEOUT));
220                let response = match read_request(&mut stream) {
221                    Ok(request) => route(request, token, &mut active_session, &command_tx),
222                    Err(error) => webdriver_error("invalid argument", error.to_string()),
223                };
224                let _ = write_response(&mut stream, response);
225            }
226            Err(_) if shutdown.load(Ordering::Acquire) => break,
227            Err(_) => continue,
228        }
229    }
230}
231
232#[derive(Debug)]
233struct HttpRequest {
234    method: String,
235    path: String,
236    body: Value,
237}
238
239fn read_request(stream: &mut TcpStream) -> io::Result<HttpRequest> {
240    let mut bytes = Vec::with_capacity(4096);
241    let header_end = loop {
242        if bytes.len() >= MAX_REQUEST_BYTES {
243            return Err(io::Error::new(
244                io::ErrorKind::InvalidData,
245                "request is too large",
246            ));
247        }
248        let mut chunk = [0; 4096];
249        let count = stream.read(&mut chunk)?;
250        if count == 0 {
251            return Err(io::Error::new(
252                io::ErrorKind::UnexpectedEof,
253                "connection closed before headers",
254            ));
255        }
256        bytes.extend_from_slice(&chunk[..count]);
257        if let Some(index) = bytes.windows(4).position(|window| window == b"\r\n\r\n") {
258            break index + 4;
259        }
260    };
261
262    let headers = std::str::from_utf8(&bytes[..header_end])
263        .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
264    let mut lines = headers.split("\r\n");
265    let mut request_line = lines.next().unwrap_or_default().split_whitespace();
266    let method = request_line.next().unwrap_or_default().to_string();
267    let path = request_line.next().unwrap_or_default().to_string();
268    if method.is_empty() || path.is_empty() {
269        return Err(io::Error::new(
270            io::ErrorKind::InvalidData,
271            "invalid request line",
272        ));
273    }
274    let content_length = lines
275        .filter_map(|line| line.split_once(':'))
276        .find(|(name, _)| name.eq_ignore_ascii_case("content-length"))
277        .map(|(_, value)| value.trim().parse::<usize>())
278        .transpose()
279        .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?
280        .unwrap_or(0);
281    if header_end + content_length > MAX_REQUEST_BYTES {
282        return Err(io::Error::new(
283            io::ErrorKind::InvalidData,
284            "request is too large",
285        ));
286    }
287    while bytes.len() < header_end + content_length {
288        let remaining = header_end + content_length - bytes.len();
289        let mut chunk = vec![0; remaining.min(4096)];
290        let count = stream.read(&mut chunk)?;
291        if count == 0 {
292            return Err(io::Error::new(
293                io::ErrorKind::UnexpectedEof,
294                "connection closed before body",
295            ));
296        }
297        bytes.extend_from_slice(&chunk[..count]);
298    }
299    let body = if content_length == 0 {
300        Value::Null
301    } else {
302        serde_json::from_slice(&bytes[header_end..header_end + content_length])
303            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?
304    };
305    Ok(HttpRequest { method, path, body })
306}
307
308fn route(
309    request: HttpRequest,
310    token: &str,
311    active_session: &mut Option<String>,
312    command_tx: &SyncSender<ControlRequest>,
313) -> Value {
314    if request.method == "GET" && request.path == "/status" {
315        return json!({"value": {
316            "ready": true,
317            "message": "Blitz debug control is ready",
318            "protocolVersion": PROTOCOL_VERSION,
319        }});
320    }
321
322    if request.method == "POST" && request.path == "/session" {
323        if active_session.is_some() {
324            return webdriver_error("session not created", "only one session is supported");
325        }
326        let supplied_token = request
327            .body
328            .pointer("/capabilities/alwaysMatch/blitz:token")
329            .and_then(Value::as_str);
330        if supplied_token != Some(token) {
331            return webdriver_error("invalid argument", "invalid blitz:token capability");
332        }
333        let session_id = match random_hex(16) {
334            Ok(value) => value,
335            Err(error) => return webdriver_error("unknown error", error.to_string()),
336        };
337        *active_session = Some(session_id.clone());
338        return json!({"value": {
339            "sessionId": session_id,
340            "capabilities": {
341                "browserName": "blitz",
342                "blitz:protocolVersion": PROTOCOL_VERSION,
343            }
344        }});
345    }
346
347    let Some((session_id, command_path)) = session_path(&request.path) else {
348        return webdriver_error("unknown command", "unknown debug-control route");
349    };
350    if active_session.as_deref() != Some(session_id) {
351        return webdriver_error("invalid session id", "session is not active");
352    }
353    if request.method == "DELETE" && command_path.is_empty() {
354        *active_session = None;
355        return json!({"value": null});
356    }
357
358    let (reply_tx, reply_rx) = mpsc::sync_channel(1);
359    let control_request = ControlRequest {
360        method: request.method,
361        path: command_path.to_string(),
362        body: request.body,
363        reply: reply_tx,
364    };
365    match command_tx.try_send(control_request) {
366        Ok(()) => {}
367        Err(TrySendError::Full(_)) => {
368            return webdriver_error("timeout", "renderer command queue is full");
369        }
370        Err(TrySendError::Disconnected(_)) => {
371            return webdriver_error("unknown error", "renderer command channel is closed");
372        }
373    }
374    match reply_rx.recv_timeout(COMMAND_TIMEOUT) {
375        Ok(ControlResponse::Success(value)) => json!({"value": value}),
376        Ok(ControlResponse::Error {
377            error,
378            message,
379            stacktrace,
380        }) => json!({"value": {
381            "error": error,
382            "message": message,
383            "stacktrace": stacktrace,
384        }}),
385        Err(RecvTimeoutError::Timeout) => webdriver_error("timeout", "renderer command timed out"),
386        Err(RecvTimeoutError::Disconnected) => {
387            webdriver_error("unknown error", "renderer response channel is closed")
388        }
389    }
390}
391
392fn session_path(path: &str) -> Option<(&str, &str)> {
393    let remainder = path.strip_prefix("/session/")?;
394    let (session_id, command) = remainder.split_once('/').unwrap_or((remainder, ""));
395    Some((session_id, command))
396}
397
398fn webdriver_error(error: &str, message: impl Into<String>) -> Value {
399    json!({"value": {
400        "error": error,
401        "message": message.into(),
402        "stacktrace": "",
403    }})
404}
405
406fn write_response(stream: &mut TcpStream, body: Value) -> io::Result<()> {
407    let bytes = serde_json::to_vec(&body).map_err(io::Error::other)?;
408    write!(
409        stream,
410        "HTTP/1.1 200 OK\r\nContent-Type: application/json; charset=utf-8\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
411        bytes.len()
412    )?;
413    stream.write_all(&bytes)
414}
415
416#[cfg(test)]
417mod tests {
418    use super::*;
419    use std::time::{SystemTime, UNIX_EPOCH};
420
421    fn descriptor_path() -> PathBuf {
422        let nonce = SystemTime::now()
423            .duration_since(UNIX_EPOCH)
424            .unwrap()
425            .as_nanos();
426        std::env::temp_dir().join(format!("blitz-debug-{nonce}.json"))
427    }
428
429    fn request(address: SocketAddr, method: &str, path: &str, body: Value) -> Value {
430        let body = if body.is_null() {
431            Vec::new()
432        } else {
433            serde_json::to_vec(&body).unwrap()
434        };
435        let mut stream = TcpStream::connect(address).unwrap();
436        write!(
437            stream,
438            "{method} {path} HTTP/1.1\r\nHost: {address}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
439            body.len()
440        )
441        .unwrap();
442        stream.write_all(&body).unwrap();
443        let mut response = Vec::new();
444        stream.read_to_end(&mut response).unwrap();
445        let body_start = response
446            .windows(4)
447            .position(|window| window == b"\r\n\r\n")
448            .unwrap()
449            + 4;
450        serde_json::from_slice(&response[body_start..]).unwrap()
451    }
452
453    fn create_session(address: SocketAddr, token: &str) -> String {
454        request(
455            address,
456            "POST",
457            "/session",
458            json!({"capabilities": {"alwaysMatch": {"blitz:token": token}}}),
459        )["value"]["sessionId"]
460            .as_str()
461            .unwrap()
462            .to_string()
463    }
464
465    #[test]
466    fn status_auth_session_command_and_reconnect() {
467        let descriptor = descriptor_path();
468        let (server, commands) = DebugServer::start(ServerConfig {
469            bind_address: (std::net::Ipv4Addr::LOCALHOST, 0).into(),
470            descriptor_path: descriptor.clone(),
471            renderer_revision: "test-revision".into(),
472        })
473        .unwrap();
474
475        let status = request(server.address(), "GET", "/status", Value::Null);
476        assert_eq!(status["value"]["ready"], true);
477        assert!(descriptor.exists());
478        #[cfg(unix)]
479        {
480            use std::os::unix::fs::PermissionsExt;
481            assert_eq!(
482                fs::metadata(&descriptor).unwrap().permissions().mode() & 0o777,
483                0o600
484            );
485        }
486
487        let rejected = request(
488            server.address(),
489            "POST",
490            "/session",
491            json!({"capabilities": {"alwaysMatch": {"blitz:token": "wrong"}}}),
492        );
493        assert_eq!(rejected["value"]["error"], "invalid argument");
494
495        let session = create_session(server.address(), server.token());
496        let address = server.address();
497        let command_path = format!("/session/{session}/blitz/getDomSnapshot");
498        let client = thread::spawn(move || request(address, "GET", &command_path, Value::Null));
499        let command = commands.recv_timeout(COMMAND_TIMEOUT).unwrap();
500        assert_eq!(command.method, "GET");
501        assert_eq!(command.path, "blitz/getDomSnapshot");
502        command
503            .respond(ControlResponse::Success(json!({"documentRevision": 7})))
504            .unwrap();
505        let response = client.join().unwrap();
506        assert_eq!(response["value"]["documentRevision"], 7);
507
508        let deleted = request(
509            server.address(),
510            "DELETE",
511            &format!("/session/{session}"),
512            Value::Null,
513        );
514        assert!(deleted["value"].is_null());
515        let second_session = create_session(server.address(), server.token());
516        assert_ne!(second_session, session);
517
518        server.shutdown();
519        assert!(!descriptor.exists());
520    }
521
522    #[test]
523    fn rejects_non_loopback_bind_address() {
524        let result = DebugServer::start(ServerConfig {
525            bind_address: ([0, 0, 0, 0], 0).into(),
526            descriptor_path: descriptor_path(),
527            renderer_revision: "test-revision".into(),
528        });
529        assert!(matches!(result, Err(error) if error.kind() == io::ErrorKind::InvalidInput));
530    }
531
532    #[test]
533    fn a_full_renderer_queue_fails_without_blocking_the_server() {
534        let (command_tx, command_rx) = mpsc::sync_channel(1);
535        let (reply_tx, _reply_rx) = mpsc::sync_channel(1);
536        command_tx
537            .send(ControlRequest {
538                method: "GET".into(),
539                path: "occupied".into(),
540                body: Value::Null,
541                reply: reply_tx,
542            })
543            .unwrap();
544        let mut active_session = Some("test-session".to_string());
545        let response = route(
546            HttpRequest {
547                method: "GET".into(),
548                path: "/session/test-session/blitz/getDomSnapshot".into(),
549                body: Value::Null,
550            },
551            "token",
552            &mut active_session,
553            &command_tx,
554        );
555        assert_eq!(response["value"]["error"], "timeout");
556        drop(command_rx);
557    }
558}
ps_blitz_debug_control/lib.rs

ps_blitz_debug_control/
lib.rs
