Vectrum
Sign In
163 lines · 6.8 KB · Rust
//! Device streams for other devices, as T3 Code passes them through its
//! environment server. A client's Device panel connects to this device's
//! remote server under `/device-hub/…` with the key; the request goes on to
//! the device hub on loopback and the connection is spliced through, so
//! video (H.264 over a WebSocket, or MJPEG) and touches flow straight
//! between the two.
//!
//! The hub never faces the network itself: it has routes that run commands.
//! Only the stream and input routes pass, and only GETs (WebSocket upgrades
//! are GETs too). Everything else a client does with devices goes through
//! typed messages on the remote connection.

use std::{
    io::{Read as _, Write as _},
    net::{Shutdown, TcpStream},
    thread,
    time::Duration,
};

use super::keys_match;

/// Where device routes start on the remote server.
pub const PREFIX: &str = "/device-hub";
/// The most a request head may take.
const MAX_HEAD: usize = 16 * 1024;

/// Looks at a new connection's request head without taking it, and says
/// whether it's for the device hub.
pub fn is_device_request(stream: &TcpStream) -> bool {
    let mut buffer = [0u8; 512];
    let deadline = std::time::Instant::now() + Duration::from_secs(5);
    loop {
        match stream.peek(&mut buffer) {
            Ok(0) => return false,
            Ok(n) => {
                let head = String::from_utf8_lossy(&buffer[..n]);
                if let Some(line) = head.split("\r\n").next().filter(|_| head.contains("\r\n")) {
                    return line.split_whitespace().nth(1).is_some_and(|path| path.starts_with(&format!("{PREFIX}/")));
                }
                if n == buffer.len() || std::time::Instant::now() > deadline {
                    return false;
                }
                thread::sleep(Duration::from_millis(10));
            }
            Err(_) => return false,
        }
    }
}

/// Whether a hub route may be reached from another device.
fn allowed(path: &str) -> bool {
    let route = path.split('?').next().unwrap_or_default();
    match route {
        // Android: H.264 out, gestures in.
        "/vendor/serve-emu/ws" => true,
        // iOS: input.
        "/vendor/serve-sim/helper/ws" => true,
        // iOS: video, `/vendor/serve-sim/helper/<udid>/stream.mjpeg`.
        _ => route
            .strip_prefix("/vendor/serve-sim/helper/")
            .and_then(|rest| rest.strip_suffix("/stream.mjpeg"))
            .is_some_and(|udid| !udid.is_empty() && !udid.contains('/')),
    }
}

fn respond(stream: &mut TcpStream, status: &str) {
    let body = status.split_once(' ').map(|(_, text)| text).unwrap_or(status);
    let _ = write!(stream, "HTTP/1.1 {status}\r\nContent-Type: text/plain\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len());
    let _ = stream.flush();
}

/// Serves one device request: checks the key and the route, then splices
/// the connection through to the hub. `hub` finds the hub's loopback
/// origin, starting it if device support is on.
pub fn serve(mut stream: TcpStream, key: &str, hub: impl FnOnce() -> Option<String>) {
    let _ = stream.set_read_timeout(Some(Duration::from_secs(10)));
    let Some(head) = read_head(&mut stream) else { return };
    let mut lines = head.split("\r\n");
    let request_line = lines.next().unwrap_or_default();
    let mut parts = request_line.split_whitespace();
    let (method, path, version) = (parts.next().unwrap_or_default(), parts.next().unwrap_or_default(), parts.next().unwrap_or("HTTP/1.1"));
    let headers: Vec<(&str, &str)> = lines
        .filter(|line| !line.is_empty())
        .filter_map(|line| line.split_once(':').map(|(name, value)| (name.trim(), value.trim())))
        .collect();
    let offered = headers
        .iter()
        .find(|(name, _)| name.eq_ignore_ascii_case("authorization"))
        .and_then(|(_, value)| value.strip_prefix("Bearer "))
        .unwrap_or_default();
    if !keys_match(offered, key) {
        // Slow down anyone guessing keys.
        thread::sleep(Duration::from_secs(1));
        return respond(&mut stream, "401 Unauthorized");
    }
    let route = path.strip_prefix(PREFIX).unwrap_or_default();
    if method != "GET" || !allowed(route) {
        return respond(&mut stream, "403 Forbidden");
    }
    let Some(origin) = hub() else { return respond(&mut stream, "503 Service Unavailable") };
    let address = origin.trim_start_matches("http://").to_string();
    let Ok(mut upstream) = TcpStream::connect(&address) else { return respond(&mut stream, "502 Bad Gateway") };
    let _ = upstream.set_nodelay(true);
    let _ = stream.set_nodelay(true);

    // The same request, for the hub: its own path, its own host, no key.
    let mut forwarded = format!("{method} {route} {version}\r\n");
    for (name, value) in &headers {
        if name.eq_ignore_ascii_case("authorization") || name.eq_ignore_ascii_case("host") {
            continue;
        }
        forwarded.push_str(&format!("{name}: {value}\r\n"));
    }
    forwarded.push_str(&format!("Host: {address}\r\n\r\n"));
    if upstream.write_all(forwarded.as_bytes()).is_err() {
        return respond(&mut stream, "502 Bad Gateway");
    }
    splice(stream, upstream);
}

/// Reads the request head, up to and including its blank line.
fn read_head(stream: &mut TcpStream) -> Option<String> {
    let mut head = Vec::new();
    let mut byte = [0u8; 1];
    while !head.ends_with(b"\r\n\r\n") {
        if head.len() >= MAX_HEAD || stream.read(&mut byte).ok()? == 0 {
            return None;
        }
        head.push(byte[0]);
    }
    Some(String::from_utf8_lossy(&head).into_owned())
}

/// Copies both ways until either side hangs up.
fn splice(client: TcpStream, hub: TcpStream) {
    let _ = client.set_read_timeout(None);
    let (Ok(mut client_read), Ok(mut hub_write)) = (client.try_clone(), hub.try_clone()) else { return };
    let upward = thread::spawn(move || {
        let _ = std::io::copy(&mut client_read, &mut hub_write);
        let _ = hub_write.shutdown(Shutdown::Both);
    });
    let (mut hub_read, mut client_write) = (hub, client);
    let _ = std::io::copy(&mut hub_read, &mut client_write);
    let _ = client_write.shutdown(Shutdown::Both);
    let _ = upward.join();
}

#[cfg(test)]
mod tests {
    use super::allowed;

    #[::core::prelude::v1::test]
    fn passes_only_stream_routes() {
        assert!(allowed("/vendor/serve-emu/ws?device=emulator-5554&frame-meta=1"));
        assert!(allowed("/vendor/serve-sim/helper/ws?device=ABC"));
        assert!(allowed("/vendor/serve-sim/helper/ABC-123/stream.mjpeg"));
        assert!(!allowed("/vendor/serve-sim/helper/a/b/stream.mjpeg"));
        assert!(!allowed("/vendor/serve-sim/api"));
        assert!(!allowed("/vendor/serve-sim/exec"));
        assert!(!allowed("/api/devices/boot"));
        assert!(!allowed("/vendor/serve-emu/api/screenshot?device=x"));
    }
}