Genarrative/apps/desktop-shell/src-tauri/src/host_bridge/protocol.rs

use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use std::sync::{Arc, Condvar, Mutex};

pub(crate) const HOST_BRIDGE_PROTOCOL: &str = "GenarrativeHostBridge";
pub(crate) const HOST_BRIDGE_VERSION: u8 = 1;
pub(crate) const HOST_BRIDGE_METHODS: [&str; 25] = [
    "host.getRuntime",
    "appearance.getColorScheme",
    "auth.requestLogin",
    "payment.request",
    "share.setTarget",
    "share.open",
    "navigation.openNativePage",
    "app.reloadWebView",
    "app.openExternalUrl",
    "app.setTitle",
    "app.setBadgeCount",
    "network.status",
    "clipboard.writeText",
    "clipboard.readText",
    "file.exportText",
    "file.importText",
    "file.importDocument",
    "file.exportImage",
    "file.importImage",
    "file.captureImage",
    "scanner.scanQrCode",
    "file.importAudio",
    "file.exportAudio",
    "haptics.impact",
    "notification.showLocal",
];
pub(crate) const HOST_BRIDGE_REQUEST_ID_MAX_LENGTH: usize = 120;
const HOST_BRIDGE_RESPONSE_CACHE_MAX: usize = 128;

#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct HostBridgeRequest {
    pub(crate) bridge: String,
    pub(crate) version: u8,
    pub(crate) id: String,
    pub(crate) method: String,
    pub(crate) payload: Option<Value>,
}

#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct HostBridgeRuntime {
    pub(crate) shell: &'static str,
    pub(crate) platform: &'static str,
    pub(crate) host_version: &'static str,
    pub(crate) bridge_version: u8,
    pub(crate) capabilities: Vec<&'static str>,
}

#[derive(Clone, Debug, Serialize)]
pub(crate) struct HostBridgeError {
    pub(crate) code: &'static str,
    pub(crate) message: String,
}

#[derive(Clone, Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct HostBridgeResponse {
    pub(crate) bridge: &'static str,
    pub(crate) version: u8,
    pub(crate) id: String,
    pub(crate) ok: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub(crate) result: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub(crate) error: Option<HostBridgeError>,
}

#[derive(Debug, Default)]
pub(crate) struct HostBridgeReplayState {
    cache: Mutex<HostBridgeReplayCache>,
}

#[derive(Debug, Default)]
struct HostBridgeReplayCache {
    order: Vec<String>,
    slots: HashMap<String, Arc<HostBridgeReplaySlot>>,
}

#[derive(Debug, Default)]
pub(crate) struct HostBridgeReplaySlot {
    response: Mutex<Option<HostBridgeResponse>>,
    ready: Condvar,
}

#[derive(Debug)]
pub(crate) enum HostBridgeReplayReservation {
    Execute(Arc<HostBridgeReplaySlot>),
    Wait(Arc<HostBridgeReplaySlot>),
}

impl HostBridgeReplayState {
    pub(crate) fn reserve(&self, request_id: &str) -> HostBridgeReplayReservation {
        let mut cache = self.cache.lock().expect("host bridge replay cache lock");
        if let Some(slot) = cache.slots.get(request_id) {
            return HostBridgeReplayReservation::Wait(slot.clone());
        }

        let slot = Arc::new(HostBridgeReplaySlot::default());
        cache.order.push(request_id.to_string());
        cache.slots.insert(request_id.to_string(), slot.clone());
        while cache.order.len() > HOST_BRIDGE_RESPONSE_CACHE_MAX {
            if let Some(oldest_request_id) = cache.order.first().cloned() {
                cache.order.remove(0);
                cache.slots.remove(&oldest_request_id);
            }
        }

        HostBridgeReplayReservation::Execute(slot)
    }

    pub(crate) fn complete(
        &self,
        slot: Arc<HostBridgeReplaySlot>,
        response: HostBridgeResponse,
    ) -> HostBridgeResponse {
        let mut stored_response = slot.response.lock().expect("host bridge replay slot lock");
        *stored_response = Some(response.clone());
        slot.ready.notify_all();
        response
    }

    pub(crate) fn wait_for_response(slot: Arc<HostBridgeReplaySlot>) -> HostBridgeResponse {
        let mut stored_response = slot.response.lock().expect("host bridge replay slot lock");
        while stored_response.is_none() {
            stored_response = slot
                .ready
                .wait(stored_response)
                .expect("host bridge replay slot wait");
        }

        stored_response
            .clone()
            .expect("host bridge replay response")
    }
}

pub(crate) fn ok(id: String, result: Value) -> HostBridgeResponse {
    HostBridgeResponse {
        bridge: HOST_BRIDGE_PROTOCOL,
        version: HOST_BRIDGE_VERSION,
        id,
        ok: true,
        result: Some(result),
        error: None,
    }
}

pub(crate) fn failed(
    id: String,
    code: &'static str,
    message: impl Into<String>,
) -> HostBridgeResponse {
    HostBridgeResponse {
        bridge: HOST_BRIDGE_PROTOCOL,
        version: HOST_BRIDGE_VERSION,
        id,
        ok: false,
        result: None,
        error: Some(HostBridgeError {
            code,
            message: message.into(),
        }),
    }
}

pub(crate) fn has_control_character(value: &str) -> bool {
    value.chars().any(|character| {
        let code_point = character as u32;
        code_point <= 31 || code_point == 127
    })
}

pub(crate) fn normalize_request_id(raw_id: &str) -> Option<String> {
    let id = raw_id.trim();
    if id.is_empty()
        || id.chars().count() > HOST_BRIDGE_REQUEST_ID_MAX_LENGTH
        || has_control_character(id)
    {
        return None;
    }

    Some(id.to_string())
}

pub(crate) fn is_host_bridge_method(method: &str) -> bool {
    HOST_BRIDGE_METHODS.contains(&method)
}

pub(crate) fn validate_request(request: &HostBridgeRequest) -> Option<HostBridgeResponse> {
    let Some(request_id) = normalize_request_id(&request.id) else {
        return Some(failed(
            "invalid".to_string(),
            "invalid_request",
            "invalid host bridge request id",
        ));
    };

    if request.bridge != HOST_BRIDGE_PROTOCOL || request.version != HOST_BRIDGE_VERSION {
        return Some(failed(
            request_id,
            "invalid_request",
            "invalid host bridge envelope",
        ));
    }

    if !is_host_bridge_method(&request.method) {
        return Some(failed(
            request_id,
            "invalid_request",
            "invalid host bridge method",
        ));
    }

    None
}

pub(crate) fn required_string_payload<'a>(
    request: &'a HostBridgeRequest,
    field: &'static str,
) -> Result<&'a str, HostBridgeResponse> {
    request
        .payload
        .as_ref()
        .and_then(|value| value.get(field))
        .and_then(Value::as_str)
        .ok_or_else(|| {
            failed(
                request.id.clone(),
                "invalid_request",
                format!("{} is required", field),
            )
        })
}

#[cfg(test)]
pub(crate) fn request(method: &str) -> HostBridgeRequest {
    HostBridgeRequest {
        bridge: HOST_BRIDGE_PROTOCOL.to_string(),
        version: HOST_BRIDGE_VERSION,
        id: "request-1".to_string(),
        method: method.to_string(),
        payload: None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn invalid_envelope_is_rejected() {
        let mut invalid = request("host.getRuntime");
        invalid.bridge = "OtherBridge".to_string();

        let response = validate_request(&invalid).expect("invalid envelope");

        assert!(!response.ok);
        assert_eq!(response.error.expect("error").code, "invalid_request");
    }

    #[test]
    fn invalid_request_id_and_unknown_method_are_rejected() {
        for id in ["", "request\n1"] {
            let mut invalid = request("share.open");
            invalid.id = id.to_string();

            let response = validate_request(&invalid).expect("invalid id");

            assert!(!response.ok);
            assert_eq!(response.id, "invalid");
            assert_eq!(response.error.expect("error").code, "invalid_request");
        }

        let mut oversized = request("share.open");
        oversized.id = "a".repeat(HOST_BRIDGE_REQUEST_ID_MAX_LENGTH + 1);
        let response = validate_request(&oversized).expect("oversized id");
        assert!(!response.ok);
        assert_eq!(response.id, "invalid");
        assert_eq!(response.error.expect("error").code, "invalid_request");

        let mut multibyte_boundary = request("host.getRuntime");
        multibyte_boundary.id = "作".repeat(HOST_BRIDGE_REQUEST_ID_MAX_LENGTH);
        assert!(validate_request(&multibyte_boundary).is_none());

        let response =
            validate_request(&request("host.runArbitraryCommand")).expect("unknown method");
        assert!(!response.ok);
        let error = response.error.expect("error");
        assert_eq!(error.code, "invalid_request");
        assert_eq!(error.message, "invalid host bridge method");
    }

    #[test]
    fn host_bridge_replay_state_reuses_first_response_for_duplicate_id() {
        let replay_state = HostBridgeReplayState::default();
        let mut side_effect_count = 0;

        let first_reservation = replay_state.reserve("request-1");
        let first_response = match first_reservation {
            HostBridgeReplayReservation::Execute(slot) => {
                side_effect_count += 1;
                replay_state.complete(slot, ok("request-1".to_string(), json!(true)))
            }
            HostBridgeReplayReservation::Wait(_) => panic!("first request must execute"),
        };
        let second_response = match replay_state.reserve("request-1") {
            HostBridgeReplayReservation::Execute(_) => panic!("duplicate request must not execute"),
            HostBridgeReplayReservation::Wait(slot) => {
                HostBridgeReplayState::wait_for_response(slot)
            }
        };

        assert_eq!(side_effect_count, 1);
        assert_eq!(second_response.ok, first_response.ok);
        assert_eq!(second_response.result, first_response.result);
    }

    #[test]
    fn host_bridge_replay_state_evicts_oldest_response_after_cache_limit() {
        let replay_state = HostBridgeReplayState::default();

        match replay_state.reserve("request-0") {
            HostBridgeReplayReservation::Execute(slot) => {
                replay_state.complete(slot, ok("request-0".to_string(), json!(0)));
            }
            HostBridgeReplayReservation::Wait(_) => panic!("first request must execute"),
        }

        for index in 1..=HOST_BRIDGE_RESPONSE_CACHE_MAX {
            let request_id = format!("request-{index}");
            match replay_state.reserve(&request_id) {
                HostBridgeReplayReservation::Execute(slot) => {
                    replay_state.complete(slot, ok(request_id, json!(index)));
                }
                HostBridgeReplayReservation::Wait(_) => panic!("new request must execute"),
            }
        }

        match replay_state.reserve("request-0") {
            HostBridgeReplayReservation::Execute(_) => {}
            HostBridgeReplayReservation::Wait(_) => {
                panic!("oldest request must be evicted after cache limit")
            }
        }
    }

    #[test]
    fn invalid_string_payload_is_rejected() {
        let mut invalid = request("clipboard.writeText");
        invalid.payload = Some(json!({ "text": 123 }));

        let response = required_string_payload(&invalid, "text").expect_err("invalid payload");

        assert!(!response.ok);
        let error = response.error.expect("error");
        assert_eq!(error.code, "invalid_request");
        assert_eq!(error.message, "text is required");
    }
}