tsafe-nativehost 1.0.5

//! tsafe Native Messaging Host — library entry point exposed for the tsafe meta-crate.
//!
//! All host logic is defined here so the meta-crate can link against it.
//! `main.rs` is a thin shim that calls [`run`].

use std::collections::HashMap;
use std::io::{self, Read, Write};

use chrono::Utc;
use serde_json::{json, Value};
use tsafe_core::{
    audit::{AuditEntry, AuditLog},
    errors::SafeError,
    profile,
    vault::Vault,
};
use zeroize::Zeroize;

struct Session {
    token: String,
    profile: String,
    password: String,
    expires_at: chrono::DateTime<Utc>,
}

impl Drop for Session {
    fn drop(&mut self) {
        self.password.zeroize();
        self.token.zeroize();
    }
}

fn message_hostname(msg: &Value) -> Option<&str> {
    msg["hostname"]
        .as_str()
        .map(str::trim)
        .filter(|s| !s.is_empty())
}

fn validated_message_hostname(msg: &Value) -> Result<&str, Value> {
    let hostname = message_hostname(msg)
        .ok_or_else(|| json!({"status": "error", "message": "missing hostname"}))?;
    if let Some(err) = browser_hostname_rpc_error(Some(hostname)) {
        return Err(err);
    }
    Ok(hostname)
}

fn browser_request_allowed(session: &Session, hostname: &str) -> bool {
    matches!(
        profile::resolve_browser_profile(hostname),
        Ok(Some(mapped_profile)) if mapped_profile == session.profile
    )
}

/// Reject structurally invalid hostnames before profile mapping (see [`profile::browser_hostname_fill_guard`]).
fn browser_hostname_rpc_error(hostname: Option<&str>) -> Option<Value> {
    let h = hostname?;
    profile::browser_hostname_fill_guard(h)
        .err()
        .map(|reason| json!({"status": "error", "message": format!("hostname rejected: {reason}")}))
}

fn parse_json_body(body: &[u8]) -> anyhow::Result<Value> {
    Ok(serde_json::from_slice(body)?)
}

/// Read one Chrome native-messaging frame: 4-byte little-endian length + UTF-8 JSON body.
pub(crate) fn read_message_from<R: Read>(mut r: R) -> anyhow::Result<Value> {
    let mut len_buf = [0u8; 4];
    r.read_exact(&mut len_buf)?;
    let len = u32::from_le_bytes(len_buf) as usize;
    if len > 1_048_576 {
        anyhow::bail!("message too large: {len}");
    }
    let mut body = vec![0u8; len];
    r.read_exact(&mut body)?;
    parse_json_body(&body)
}

fn read_message() -> anyhow::Result<Value> {
    read_message_from(io::stdin().lock())
}

fn write_message(resp: &Value) -> anyhow::Result<()> {
    let bytes = serde_json::to_vec(resp)?;
    let len = bytes.len() as u32;
    io::stdout().write_all(&len.to_le_bytes())?;
    io::stdout().write_all(&bytes)?;
    io::stdout().flush()?;
    Ok(())
}

/// Check whether `token` matches the active session and has not expired.
/// Refreshes the TTL on success; drops the session on expiry.
fn check_session(session: &mut Option<Session>, token: &str) -> bool {
    if let Some(ref mut s) = session {
        if s.token == token {
            if Utc::now() < s.expires_at {
                s.expires_at = Utc::now() + chrono::Duration::minutes(5);
                return true;
            } else {
                *session = None;
            }
        }
    }
    false
}

fn append_browser_audit(profile_name: &str, operation: &str, key: Option<&str>) {
    AuditLog::new(&profile::audit_log_path(profile_name))
        .append(&AuditEntry::success(profile_name, operation, key))
        .ok();
}

pub(crate) fn dispatch(session: &mut Option<Session>, msg: Value) -> Value {
    let command = msg["command"].as_str().unwrap_or("");
    match command {
        "unlock" => {
            let profile_name = msg["profile"].as_str().unwrap_or("default").to_string();
            if profile::validate_profile_name(&profile_name).is_err() {
                return json!({"status": "error", "message": "invalid profile name"});
            }
            let password = match msg["password"].as_str() {
                Some(p) => p.to_string(),
                None => return json!({"status": "error", "message": "missing password"}),
            };
            let vault_path = profile::vault_path(&profile_name);
            match Vault::open(&vault_path, password.as_bytes()) {
                Ok(_) => {
                    let token = uuid::Uuid::new_v4().to_string();
                    let expires_at = Utc::now() + chrono::Duration::minutes(5);
                    *session = Some(Session {
                        token: token.clone(),
                        profile: profile_name.clone(),
                        password,
                        expires_at,
                    });
                    json!({
                        "status": "ok",
                        "session_token": token,
                        "profile": profile_name,
                        "expires_at": expires_at.to_rfc3339(),
                    })
                }
                Err(SafeError::DecryptionFailed) => {
                    json!({"status": "error", "message": "wrong password"})
                }
                Err(SafeError::VaultNotFound { .. }) => {
                    json!({"status": "error", "message": format!("no vault for profile '{profile_name}'")})
                }
                Err(e) => json!({"status": "error", "message": e.to_string()}),
            }
        }

        "lock" => {
            *session = None;
            json!({"status": "ok"})
        }

        "list_logins" => {
            let token = msg["session_token"].as_str().unwrap_or("");
            if !check_session(session, token) {
                return json!({"status": "error", "message": "session expired"});
            }
            let sess = session.as_ref().unwrap();
            let hostname = match validated_message_hostname(&msg) {
                Ok(hostname) => hostname,
                Err(err) => return err,
            };
            if !browser_request_allowed(sess, hostname) {
                // Before returning a generic "not mapped" error, check whether the
                // hostname is a typosquat of a registered domain (phishing guard).
                let profiles = profile::load_browser_profiles().unwrap_or_default();
                if let Some(m) = profile::lookalike_check(hostname, &profiles) {
                    return json!({
                        "status": "phishing_warning",
                        "registered": m.registered,
                    });
                }
                return json!({"status": "error", "message": "domain is not mapped to the unlocked profile"});
            }
            let vault_path = profile::vault_path(&sess.profile);
            match Vault::open(&vault_path, sess.password.as_bytes()) {
                Ok(vault) => {
                    let mut logins: Vec<Value> = vault
                        .file()
                        .secrets
                        .iter()
                        .filter(|(_, e)| {
                            // Exclude alias entries from the logins list.
                            e.tags.get("type").map(|t| t != "alias").unwrap_or(true)
                        })
                        .map(|(k, e)| {
                            let pinned = e.tags.get("pinned").map(|v| v == "true").unwrap_or(false);
                            json!({"key": k, "pinned": pinned})
                        })
                        .collect();
                    // Pinned first, then alphabetical.
                    logins.sort_by(|a, b| {
                        let ap = a["pinned"].as_bool().unwrap_or(false);
                        let bp = b["pinned"].as_bool().unwrap_or(false);
                        bp.cmp(&ap).then_with(|| {
                            a["key"]
                                .as_str()
                                .unwrap_or("")
                                .cmp(b["key"].as_str().unwrap_or(""))
                        })
                    });
                    append_browser_audit(&sess.profile, "browser-list", None);
                    json!({"status": "ok", "logins": logins})
                }
                Err(e) => json!({"status": "error", "message": e.to_string()}),
            }
        }

        "get_login" => {
            let token = msg["session_token"].as_str().unwrap_or("");
            if !check_session(session, token) {
                return json!({"status": "error", "message": "session expired"});
            }
            let key = match msg["key"].as_str() {
                Some(k) => k.to_string(),
                None => return json!({"status": "error", "message": "missing key"}),
            };
            let sess = session.as_ref().unwrap();
            let hostname = match validated_message_hostname(&msg) {
                Ok(hostname) => hostname,
                Err(err) => return err,
            };
            if !browser_request_allowed(sess, hostname) {
                return json!({"status": "error", "message": "domain is not mapped to the unlocked profile"});
            }
            let vault_path = profile::vault_path(&sess.profile);
            match Vault::open(&vault_path, sess.password.as_bytes()) {
                Ok(vault) => match vault.get(&key) {
                    Ok(value) => {
                        append_browser_audit(&sess.profile, "browser-get", Some(&key));
                        json!({"status": "ok", "value": value.as_ref() as &str})
                    }
                    Err(_) => json!({"status": "error", "message": "key not found"}),
                },
                Err(e) => json!({"status": "error", "message": e.to_string()}),
            }
        }

        "save_login" => {
            let token = msg["session_token"].as_str().unwrap_or("");
            if !check_session(session, token) {
                return json!({"status": "error", "message": "session expired"});
            }
            let key = match msg["key"].as_str() {
                Some(k) => k.to_string(),
                None => return json!({"status": "error", "message": "missing key"}),
            };
            let value = match msg["value"].as_str() {
                Some(v) => v.to_string(),
                None => return json!({"status": "error", "message": "missing value"}),
            };
            let sess = session.as_ref().unwrap();
            let hostname = match validated_message_hostname(&msg) {
                Ok(hostname) => hostname,
                Err(err) => return err,
            };
            if !browser_request_allowed(sess, hostname) {
                return json!({"status": "error", "message": "domain is not mapped to the unlocked profile"});
            }
            let vault_path = profile::vault_path(&sess.profile);
            match Vault::open(&vault_path, sess.password.as_bytes()) {
                Ok(mut vault) => match vault.set(&key, &value, HashMap::new()) {
                    Ok(()) => {
                        append_browser_audit(&sess.profile, "browser-save", Some(&key));
                        json!({"status": "ok"})
                    }
                    Err(e) => json!({"status": "error", "message": e.to_string()}),
                },
                Err(e) => json!({"status": "error", "message": e.to_string()}),
            }
        }

        other => json!({"status": "error", "message": format!("unknown command: {other}")}),
    }
}

fn message_loop() {
    let mut session: Option<Session> = None;
    loop {
        match read_message() {
            Ok(msg) => {
                // Echo the request `id` back on the response so the
                // extension's NativePortClient can route concurrent
                // requests to the right awaiter (popup + content-script
                // can interleave on the same port).
                let id = msg.get("id").cloned();
                let mut resp = dispatch(&mut session, msg);
                if let Some(id) = id {
                    if let Some(obj) = resp.as_object_mut() {
                        obj.insert("id".to_string(), id);
                    }
                }
                if let Err(e) = write_message(&resp) {
                    eprintln!("nativehost: write error: {e}");
                    break;
                }
            }
            Err(e) => {
                // EOF or broken pipe — browser closed the connection, exit cleanly.
                let s = e.to_string();
                if s.contains("failed to fill whole buffer")
                    || s.contains("unexpected end")
                    || s.contains("os error")
                {
                    break;
                }
                eprintln!("nativehost: read error: {e}");
                break;
            }
        }
    }
}

/// Launch the tsafe native messaging host.
///
/// Chrome Native Messaging uses a binary frame protocol. Switches stdin/stdout
/// to binary mode on Windows before any I/O, then enters the message loop.
pub fn run() {
    // Windows text mode mutates stdin/stdout bytes, so switch both handles
    // to binary mode before any host I/O.
    #[cfg(target_os = "windows")]
    {
        unsafe extern "C" {
            fn _setmode(fd: std::ffi::c_int, mode: std::ffi::c_int) -> std::ffi::c_int;
        }

        unsafe {
            _setmode(0, 0x8000);
            _setmode(1, 0x8000);
        }
    }

    message_loop();
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;

    use tsafe_core::{profile, vault::Vault};

    fn framed_json(bytes: &[u8]) -> Vec<u8> {
        let mut v = Vec::with_capacity(4 + bytes.len());
        v.extend_from_slice(&(bytes.len() as u32).to_le_bytes());
        v.extend_from_slice(bytes);
        v
    }

    #[test]
    fn read_message_from_accepts_valid_length_prefixed_json() {
        let payload = br#"{"command":"lock"}"#;
        let v = read_message_from(Cursor::new(framed_json(payload))).unwrap();
        assert_eq!(v["command"], "lock");
    }

    #[test]
    fn read_message_from_rejects_oversized_length_without_reading_body() {
        let len = (1_048_576 + 1) as u32;
        let err = read_message_from(Cursor::new(len.to_le_bytes().to_vec())).unwrap_err();
        assert!(
            err.to_string().contains("message too large"),
            "unexpected: {err}"
        );
    }

    #[test]
    fn read_message_from_errors_on_truncated_body() {
        let mut buf = framed_json(br"{}");
        buf.truncate(4 + 1); // length says 2, only 1 byte of body
        let err = read_message_from(Cursor::new(buf)).unwrap_err();
        let s = err.to_string().to_lowercase();
        assert!(
            s.contains("failed to fill") || s.contains("unexpected end"),
            "unexpected: {err}"
        );
    }

    #[test]
    fn read_message_from_errors_on_eof_before_length() {
        let err = read_message_from(Cursor::new([])).unwrap_err();
        let s = err.to_string().to_lowercase();
        assert!(
            s.contains("failed to fill") || s.contains("unexpected end"),
            "unexpected: {err}"
        );
    }

    #[test]
    fn read_message_from_empty_body_fails_json_parse() {
        let mut v = vec![0u8; 4];
        v.copy_from_slice(&0u32.to_le_bytes());
        let err = read_message_from(Cursor::new(v)).unwrap_err();
        let s = err.to_string().to_lowercase();
        assert!(
            s.contains("eof") || s.contains("expected") || s.contains("parse"),
            "unexpected: {err}"
        );
    }

    #[test]
    fn parse_json_body_rejects_malformed_json() {
        let err = parse_json_body(br"{not json").unwrap_err();
        let msg = err.to_string();
        assert!(
            msg.contains("key must be a string") || msg.contains("EOF") || msg.contains("expected"),
            "unexpected error: {msg}"
        );
    }

    #[test]
    fn parse_json_body_rejects_truncated_payload() {
        let err = parse_json_body(r#"{"command":"#.as_bytes()).unwrap_err();
        let lower = err.to_string().to_lowercase();
        assert!(
            lower.contains("eof") || lower.contains("unexpected"),
            "unexpected error: {err}"
        );
    }

    #[test]
    fn parse_json_body_accepts_minimal_object() {
        let v = parse_json_body(br#"{"command":"lock"}"#).unwrap();
        assert_eq!(v["command"], "lock");
    }

    #[test]
    fn dispatch_unknown_command_returns_error() {
        let mut session = None;
        let r = dispatch(&mut session, json!({"command": "not-a-real-command"}));
        assert_eq!(r["status"], "error");
        assert!(r["message"]
            .as_str()
            .unwrap_or("")
            .contains("unknown command"));
    }

    #[test]
    fn dispatch_unlock_rejects_invalid_profile_name() {
        let mut session = None;
        let r = dispatch(
            &mut session,
            json!({"command": "unlock", "profile": "evil/../../x", "password": "x"}),
        );
        assert_eq!(r["status"], "error");
        assert_eq!(r["message"], "invalid profile name");
        assert!(session.is_none());
    }

    #[test]
    fn dispatch_unlock_requires_password_field() {
        let mut session = None;
        let r = dispatch(
            &mut session,
            json!({"command": "unlock", "profile": "default"}),
        );
        assert_eq!(r["status"], "error");
        assert_eq!(r["message"], "missing password");
    }

    #[test]
    fn dispatch_list_logins_without_session_returns_expired() {
        let mut session = None;
        let r = dispatch(
            &mut session,
            json!({"command": "list_logins", "session_token": "nope"}),
        );
        assert_eq!(r["status"], "error");
        assert_eq!(r["message"], "session expired");
    }

    #[test]
    fn message_hostname_trims_and_skips_empty() {
        assert_eq!(message_hostname(&json!({"hostname": "  "})), None);
        assert_eq!(
            message_hostname(&json!({"hostname": " example.com "})),
            Some("example.com")
        );
    }

    fn unlock_default_session(dir: &std::path::Path) -> (String, Option<Session>) {
        temp_env::with_var("TSAFE_VAULT_DIR", Some(dir.to_str().unwrap()), || {
            let vault_path = profile::vault_path("default");
            std::fs::create_dir_all(vault_path.parent().unwrap()).unwrap();
            let _ = std::fs::remove_file(&vault_path);
            Vault::create(&vault_path, b"pw").unwrap();
            let mut session = None;
            let unlock = dispatch(
                &mut session,
                json!({"command": "unlock", "profile": "default", "password": "pw"}),
            );
            assert_eq!(unlock["status"], "ok", "{unlock:?}");
            let token = unlock["session_token"].as_str().unwrap().to_string();
            (token, session)
        })
    }

    #[test]
    fn dispatch_browser_paths_reject_malformed_hostname() {
        let dir = tempfile::tempdir().unwrap();
        let (token, mut session) = unlock_default_session(dir.path());
        let bad_host = (0..14)
            .map(|i| format!("l{i}"))
            .collect::<Vec<_>>()
            .join(".");

        temp_env::with_var(
            "TSAFE_VAULT_DIR",
            Some(dir.path().to_str().unwrap()),
            || {
                let r = dispatch(
                    &mut session,
                    json!({
                        "command": "list_logins",
                        "session_token": token,
                        "hostname": bad_host,
                    }),
                );
                assert_eq!(r["status"], "error");
                let msg = r["message"].as_str().unwrap();
                assert!(
                    msg.contains("hostname rejected") && msg.contains("too many"),
                    "unexpected message: {msg}"
                );

                let r = dispatch(
                    &mut session,
                    json!({
                        "command": "get_login",
                        "session_token": token,
                        "hostname": bad_host,
                        "key": "any",
                    }),
                );
                assert_eq!(r["status"], "error");
                assert!(r["message"].as_str().unwrap().contains("hostname rejected"));

                let r = dispatch(
                    &mut session,
                    json!({
                        "command": "save_login",
                        "session_token": token,
                        "hostname": bad_host,
                        "key": "k",
                        "value": "v",
                    }),
                );
                assert_eq!(r["status"], "error");
                assert!(r["message"].as_str().unwrap().contains("hostname rejected"));
            },
        );
    }

    #[test]
    fn dispatch_list_logins_returns_phishing_warning_for_lookalike_hostname() {
        let dir = tempfile::tempdir().unwrap();
        let (token, mut session) = unlock_default_session(dir.path());

        // Write a browser-profiles.json mapping paypal.com to "default".
        let profiles_path = dir.path().join("browser-profiles.json");
        std::fs::write(&profiles_path, r#"{"paypal.com":"default"}"#).unwrap();

        temp_env::with_var(
            "TSAFE_VAULT_DIR",
            Some(dir.path().to_str().unwrap()),
            || {
                // paypa1.com is 1 edit away from paypal.com → phishing_warning
                let r = dispatch(
                    &mut session,
                    json!({
                        "command": "list_logins",
                        "session_token": token,
                        "hostname": "paypa1.com",
                    }),
                );
                assert_eq!(
                    r["status"].as_str().unwrap(),
                    "phishing_warning",
                    "expected phishing_warning, got: {r:?}"
                );
                assert_eq!(r["registered"].as_str().unwrap(), "paypal.com");
            },
        );
    }

    #[test]
    fn dispatch_browser_paths_require_hostname() {
        let dir = tempfile::tempdir().unwrap();
        let (token, mut session) = unlock_default_session(dir.path());

        temp_env::with_var(
            "TSAFE_VAULT_DIR",
            Some(dir.path().to_str().unwrap()),
            || {
                let r = dispatch(
                    &mut session,
                    json!({"command": "list_logins", "session_token": token}),
                );
                assert_eq!(r["status"], "error");
                assert_eq!(r["message"], "missing hostname");

                let r = dispatch(
                    &mut session,
                    json!({
                        "command": "get_login",
                        "session_token": token,
                        "key": "any",
                    }),
                );
                assert_eq!(r["status"], "error");
                assert_eq!(r["message"], "missing hostname");

                let r = dispatch(
                    &mut session,
                    json!({
                        "command": "save_login",
                        "session_token": token,
                        "key": "k",
                        "value": "v",
                    }),
                );
                assert_eq!(r["status"], "error");
                assert_eq!(r["message"], "missing hostname");
            },
        );
    }

    #[test]
    fn dispatch_preserves_session_across_messages_for_id_echo_and_continuity() {
        // Regression for the connectNative migration: under the prior
        // sendNativeMessage model, every request spawned a fresh process
        // and Session was lost. The fix is the persistent port + this
        // host's existing in-process loop. This test verifies the
        // dispatch contract that makes the persistent port work:
        //   1. session lives across dispatches on the same `&mut Option<Session>`
        //   2. the `id` field on the request is echoed back on the response
        //      (added in the main loop wrapper, NOT in dispatch — but the
        //       wrapper relies on `dispatch` returning a JSON object so
        //       insertion can happen; this test pins that invariant)
        let dir = tempfile::tempdir().unwrap();
        let (token, mut session) = unlock_default_session(dir.path());

        temp_env::with_var(
            "TSAFE_VAULT_DIR",
            Some(dir.path().to_str().unwrap()),
            || {
                let profiles_path = dir.path().join("browser-profiles.json");
                std::fs::write(&profiles_path, r#"{"github.com":"default"}"#).unwrap();

                // First call uses the session from unlock — would have
                // returned "session expired" under the broken one-shot
                // model. Under the fix, it succeeds.
                let resp1 = dispatch(
                    &mut session,
                    json!({
                        "command": "list_logins",
                        "session_token": token,
                        "hostname": "github.com",
                        "id": 42,
                    }),
                );
                assert_eq!(resp1["status"], "ok");
                // The dispatch return MUST be a JSON object so the main
                // loop can insert the echoed id. If a future refactor
                // returns a non-object value, the id-echo wrapper fails
                // silently and the extension's pending Map never resolves.
                assert!(
                    resp1.is_object(),
                    "dispatch must return a JSON object so the main loop can echo `id`"
                );

                // Second call on the SAME session — proves continuity.
                let resp2 = dispatch(
                    &mut session,
                    json!({
                        "command": "list_logins",
                        "session_token": token,
                        "hostname": "github.com",
                        "id": 43,
                    }),
                );
                assert_eq!(
                    resp2["status"], "ok",
                    "session must survive dispatch-to-dispatch"
                );
            },
        );
    }

    #[test]
    fn browser_reads_append_audit_entries() {
        let dir = tempfile::tempdir().unwrap();
        let (token, mut session) = unlock_default_session(dir.path());

        temp_env::with_var(
            "TSAFE_VAULT_DIR",
            Some(dir.path().to_str().unwrap()),
            || {
                let profiles_path = dir.path().join("browser-profiles.json");
                std::fs::write(&profiles_path, r#"{"github.com":"default"}"#).unwrap();

                let vault_path = profile::vault_path("default");
                let mut vault = Vault::open(&vault_path, b"pw").unwrap();
                vault
                    .set("github.com/octocat", "secret-value", HashMap::new())
                    .unwrap();
                drop(vault);

                let list_resp = dispatch(
                    &mut session,
                    json!({
                        "command": "list_logins",
                        "session_token": token,
                        "hostname": "github.com",
                    }),
                );
                assert_eq!(list_resp["status"], "ok");

                let get_resp = dispatch(
                    &mut session,
                    json!({
                        "command": "get_login",
                        "session_token": token,
                        "hostname": "github.com",
                        "key": "github.com/octocat",
                    }),
                );
                assert_eq!(get_resp["status"], "ok");

                let audit_path = profile::audit_log_path("default");
                let content = std::fs::read_to_string(audit_path).unwrap();
                assert!(content.contains("\"operation\":\"browser-list\""));
                assert!(content.contains("\"operation\":\"browser-get\""));
                assert!(content.contains("\"key\":\"github.com/octocat\""));
            },
        );
    }
}