host_extensions/first_party/

mesh_runtime.rs

use std::collections::{HashMap, HashSet};
use std::sync::{
    atomic::{AtomicU64, Ordering},
    Mutex,
};
use std::time::{SystemTime, UNIX_EPOCH};

use crate::{
    events::{
        MeshPrivateControlPayload, MeshPrivateReceiptPayload, MeshReplyPayload,
        MESH_PRIVATE_CONTROL, MESH_PRIVATE_RECEIPT, MESH_REPLY,
    },
    executor_contract::{
        MeshControlEnvelope, MeshObjectPolicy, MeshObjectReadReason, MeshObjectResult,
        MeshPrivateObjectRef, MeshStatusResult,
    },
    HostPushEvent,
};

/// Generate a random 32-character hex nonce for session-unique request IDs.
///
/// 16 bytes (128 bits of entropy) keeps cross-run request ID collisions
/// negligible even over very large numbers of sequential restarts.
///
/// # Fallback
///
/// When `getrandom` fails (e.g. early WASM boot without the `js` feature),
/// falls back to time + stack address + monotonic counter. This is NOT
/// cryptographically secure but is sufficient to prevent request ID reuse
/// across sequential restarts. On all supported platforms (native + WASM
/// with `js` feature), `getrandom` should always succeed and the fallback
/// should be unreachable.
pub fn generate_session_nonce() -> String {
    use std::sync::atomic::{AtomicU64, Ordering};
    static FALLBACK_CTR: AtomicU64 = AtomicU64::new(0);

    let mut bytes = [0u8; 16];
    if getrandom::getrandom(&mut bytes).is_err() {
        log::warn!("getrandom failed — using fallback nonce (not cryptographically secure)");
        let t = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_nanos()
            .to_le_bytes();
        let ctr = FALLBACK_CTR.fetch_add(1, Ordering::Relaxed).to_le_bytes();
        bytes[..8].copy_from_slice(&t[..8]);
        bytes[8..].copy_from_slice(&ctr);
    }
    bytes.iter().map(|b| format!("{b:02x}")).collect()
}

pub trait MeshRuntime: Send + Sync + 'static {
    fn publish(&self, topic: &str, data_base64: &str) -> bool;
    fn subscribe(&self, topic: &str) -> bool;
    fn put_object(&self, path: &str, data_base64: &str, policy: MeshObjectPolicy) -> bool;
    fn get_object_result(&self, path: &str) -> MeshObjectResult;
    fn get_object(&self, path: &str) -> Option<String> {
        self.get_object_result(path).compat_data()
    }
    fn request(&self, path: &str, data_base64: &str) -> String;
    fn respond(&self, request_id: &str, data_base64: &str) -> bool;
    fn status(&self) -> MeshStatusResult;
    fn supports_private(&self) -> bool {
        false
    }
    fn put_private_object(
        &self,
        _data_base64: &str,
        _policy: MeshObjectPolicy,
    ) -> MeshPrivateObjectRef {
        MeshPrivateObjectRef {
            handle: String::new(),
            capability: String::new(),
            expires_at_ms: None,
        }
    }
    fn get_private_object_result(&self, _handle: &str, _capability: &str) -> MeshObjectResult {
        MeshObjectResult::unavailable(MeshObjectReadReason::PolicyDenied, None)
    }
    fn request_private(&self, _handle: &str, _capability: &str, _data_base64: &str) -> String {
        String::new()
    }
    fn publish_private_control(&self, _capability: &str, _envelope: MeshControlEnvelope) -> bool {
        false
    }
    fn subscribe_private_control(&self, _capability: &str) -> bool {
        false
    }
    fn publish_private_receipt(&self, _capability: &str, _data_base64: &str) -> bool {
        false
    }
    fn subscribe_private_receipt(&self, _capability: &str) -> bool {
        false
    }
    fn drain_events(&self) -> Vec<HostPushEvent> {
        Vec::new()
    }
}

#[derive(Debug)]
pub struct InMemoryMeshRuntime {
    session_nonce: String,
    next_request_id: AtomicU64,
    next_private_id: AtomicU64,
    state: Mutex<InMemoryMeshState>,
}

#[derive(Debug, Default)]
struct InMemoryMeshState {
    subscriptions: HashSet<String>,
    objects: HashMap<String, StoredMeshObject>,
    private_objects: HashMap<String, StoredPrivateMeshObject>,
    private_control_subscriptions: HashSet<String>,
    private_receipt_subscriptions: HashSet<String>,
    pending_queries: HashSet<String>,
    events: Vec<HostPushEvent>,
}

// ── Wire types for statement-store transport ────────────────────────────

pub const MESH_QUERY_CHANNEL: &str = "mesh/query";
pub const MESH_REPLY_CHANNEL: &str = "mesh/reply";
pub const MESH_PRIVATE_QUERY_CHANNEL: &str = "mesh/private/query";
pub const MESH_PRIVATE_CONTROL_CHANNEL: &str = "mesh/private/control";
pub const MESH_PRIVATE_RECEIPT_CHANNEL: &str = "mesh/private/receipt";

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct MeshWireQuery {
    pub request_id: String,
    pub path: String,
    pub data_base64: String,
}

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct MeshPrivateWireQuery {
    pub request_id: String,
    pub handle: String,
    pub capability: String,
    pub data_base64: String,
}

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct MeshWireReply {
    pub request_id: String,
    pub data_base64: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reason: Option<MeshObjectReadReason>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub expires_at_ms: Option<u64>,
}

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct MeshPrivateControlWireMessage {
    pub capability: String,
    pub envelope: MeshControlEnvelope,
}

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct MeshPrivateReceiptWireMessage {
    pub capability: String,
    pub data_base64: String,
}

#[derive(Debug, Clone)]
pub struct StoredMeshObject {
    pub data_base64: String,
    pub expires_at_ms: Option<u64>,
    pub suppress_previews: bool,
}

#[derive(Debug, Clone)]
pub struct StoredPrivateMeshObject {
    pub capability: String,
    pub object: StoredMeshObject,
}

impl StoredMeshObject {
    pub fn from_put(data_base64: &str, policy: MeshObjectPolicy) -> Self {
        Self {
            data_base64: data_base64.to_string(),
            expires_at_ms: policy.expires_at_ms,
            suppress_previews: policy.suppress_previews.unwrap_or(false),
        }
    }

    pub fn as_result(&self, now_ms: u64) -> MeshObjectResult {
        // These runtimes do not synthesize previews, but they persist the
        // effective policy so compliant hosts can honor suppression without
        // changing the object-store contract.
        let _preview_suppressed = self.suppress_previews;
        if self
            .expires_at_ms
            .is_some_and(|expires_at_ms| expires_at_ms <= now_ms)
        {
            return MeshObjectResult::unavailable(
                MeshObjectReadReason::Expired,
                self.expires_at_ms,
            );
        }

        MeshObjectResult::found(self.data_base64.clone(), self.expires_at_ms)
    }
}

pub fn private_object_result(
    entry: Option<&StoredPrivateMeshObject>,
    capability: &str,
    now_ms: u64,
) -> MeshObjectResult {
    let Some(entry) = entry else {
        return MeshObjectResult::unavailable(MeshObjectReadReason::PolicyDenied, None);
    };
    if entry.capability != capability {
        return MeshObjectResult::unavailable(MeshObjectReadReason::PolicyDenied, None);
    }
    entry.object.as_result(now_ms)
}

pub fn current_time_ms() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis()
        .min(u128::from(u64::MAX)) as u64
}

impl Default for InMemoryMeshRuntime {
    fn default() -> Self {
        Self::new()
    }
}

impl InMemoryMeshRuntime {
    pub fn new() -> Self {
        Self {
            session_nonce: generate_session_nonce(),
            next_request_id: AtomicU64::new(1),
            next_private_id: AtomicU64::new(1),
            state: Mutex::new(InMemoryMeshState::default()),
        }
    }
}

impl MeshRuntime for InMemoryMeshRuntime {
    fn publish(&self, _topic: &str, _data_base64: &str) -> bool {
        true
    }

    fn subscribe(&self, topic: &str) -> bool {
        self.state
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .subscriptions
            .insert(topic.to_string());
        true
    }

    fn put_object(&self, path: &str, data_base64: &str, policy: MeshObjectPolicy) -> bool {
        self.state
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .objects
            .insert(
                path.to_string(),
                StoredMeshObject::from_put(data_base64, policy),
            );
        true
    }

    fn get_object_result(&self, path: &str) -> MeshObjectResult {
        self.state
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .objects
            .get(path)
            .map(|object| object.as_result(current_time_ms()))
            .unwrap_or_else(|| MeshObjectResult::unavailable(MeshObjectReadReason::NotFound, None))
    }

    fn request(&self, _path: &str, _data_base64: &str) -> String {
        let request_id = format!(
            "mesh-req-{}-{}",
            self.session_nonce,
            self.next_request_id.fetch_add(1, Ordering::Relaxed)
        );
        let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner());
        state.pending_queries.insert(request_id.clone());
        if state
            .objects
            .get(_path)
            .map(|object| object.as_result(current_time_ms()).data_base64.is_some())
            .unwrap_or(false)
        {
            // Known gap: cache-hit path removes pending without emitting
            // meshReply, unlike the loopback runtime. Deferred — this is
            // pre-existing InMemoryMeshRuntime behavior, not part of #117.
            state.pending_queries.remove(&request_id);
        }
        request_id
    }

    /// Marks a pending query as answered. Returns `true` if the request was pending.
    ///
    /// NOTE: Unlike `LoopbackMeshRuntime` and `StatementStoreMeshRuntime`,
    /// this implementation does NOT emit a `meshReply` event and ignores
    /// `data_base64`. This is a known asymmetry — the InMemoryMeshRuntime
    /// is for unit testing, not production mesh transport.
    fn respond(&self, request_id: &str, _data_base64: &str) -> bool {
        self.state
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .pending_queries
            .remove(request_id)
    }

    fn status(&self) -> MeshStatusResult {
        let state = self.state.lock().unwrap_or_else(|e| e.into_inner());
        MeshStatusResult {
            health: "healthy".to_string(),
            transport: "custom".to_string(),
            pending_publishes: 0,
            pending_queries: state.pending_queries.len() as u64,
            last_error: None,
        }
    }

    fn supports_private(&self) -> bool {
        true
    }

    fn put_private_object(
        &self,
        data_base64: &str,
        policy: MeshObjectPolicy,
    ) -> MeshPrivateObjectRef {
        let id = self.next_private_id.fetch_add(1, Ordering::Relaxed);
        let handle = format!("mesh-private-handle-{id}");
        let capability = format!("mesh-private-capability-{id}");
        let expires_at_ms = policy.expires_at_ms;
        self.state
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .private_objects
            .insert(
                handle.clone(),
                StoredPrivateMeshObject {
                    capability: capability.clone(),
                    object: StoredMeshObject::from_put(data_base64, policy),
                },
            );
        MeshPrivateObjectRef {
            handle,
            capability,
            expires_at_ms,
        }
    }

    fn get_private_object_result(&self, handle: &str, capability: &str) -> MeshObjectResult {
        let state = self.state.lock().unwrap_or_else(|e| e.into_inner());
        private_object_result(
            state.private_objects.get(handle),
            capability,
            current_time_ms(),
        )
    }

    fn request_private(&self, handle: &str, capability: &str, _data_base64: &str) -> String {
        let request_id = format!(
            "mesh-private-req-{}",
            self.next_request_id.fetch_add(1, Ordering::Relaxed)
        );
        let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner());
        state.pending_queries.insert(request_id.clone());
        let result = private_object_result(
            state.private_objects.get(handle),
            capability,
            current_time_ms(),
        );
        state.pending_queries.remove(&request_id);
        let payload = MeshReplyPayload {
            request_id: request_id.clone(),
            data_base64: result.data_base64,
            reason: result.reason,
            expires_at_ms: result.expires_at_ms,
            author: None,
        };
        if let Ok(payload_json) = serde_json::to_string(&payload) {
            state.events.push(HostPushEvent {
                event: MESH_REPLY.to_string(),
                payload_json,
            });
        }
        request_id
    }

    fn publish_private_control(&self, capability: &str, envelope: MeshControlEnvelope) -> bool {
        let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner());
        if state.private_control_subscriptions.contains(capability) {
            let payload = MeshPrivateControlPayload {
                capability: capability.to_string(),
                envelope,
                author: None,
            };
            if let Ok(payload_json) = serde_json::to_string(&payload) {
                state.events.push(HostPushEvent {
                    event: MESH_PRIVATE_CONTROL.to_string(),
                    payload_json,
                });
            }
        }
        true
    }

    fn subscribe_private_control(&self, capability: &str) -> bool {
        self.state
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .private_control_subscriptions
            .insert(capability.to_string());
        true
    }

    fn publish_private_receipt(&self, capability: &str, data_base64: &str) -> bool {
        let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner());
        if state.private_receipt_subscriptions.contains(capability) {
            let payload = MeshPrivateReceiptPayload {
                capability: capability.to_string(),
                data_base64: data_base64.to_string(),
                author: None,
            };
            if let Ok(payload_json) = serde_json::to_string(&payload) {
                state.events.push(HostPushEvent {
                    event: MESH_PRIVATE_RECEIPT.to_string(),
                    payload_json,
                });
            }
        }
        true
    }

    fn subscribe_private_receipt(&self, capability: &str) -> bool {
        self.state
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .private_receipt_subscriptions
            .insert(capability.to_string());
        true
    }

    fn drain_events(&self) -> Vec<HostPushEvent> {
        let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner());
        state.events.drain(..).collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::executor_contract::MeshControlMode;

    #[test]
    fn in_memory_runtime_publish_without_subscription_emits_no_events() {
        let runtime = InMemoryMeshRuntime::new();
        // Publish without subscribing first
        assert!(runtime.publish_private_control(
            "unsubscribed-capability",
            MeshControlEnvelope {
                mode: MeshControlMode::Encrypted,
                data_base64: "AQID".into(),
            },
        ));
        assert!(runtime.publish_private_receipt("unsubscribed-capability", "AQID"));
        // No events should have been emitted
        assert!(runtime.drain_events().is_empty());
    }

    #[test]
    fn in_memory_runtime_returns_structured_object_results() {
        let runtime = InMemoryMeshRuntime::new();
        assert!(runtime.put_object(
            "mesh/object/1",
            "AQID",
            MeshObjectPolicy {
                expires_at_ms: Some(u64::MAX),
                suppress_previews: Some(true),
            },
        ));

        assert_eq!(
            runtime.get_object_result("mesh/object/1"),
            MeshObjectResult::found("AQID".into(), Some(u64::MAX))
        );
        assert_eq!(runtime.get_object("mesh/object/1").as_deref(), Some("AQID"));
        assert!(runtime
            .state
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .objects
            .get("mesh/object/1")
            .map(|object| object.suppress_previews)
            .unwrap_or(false));
    }

    #[test]
    fn in_memory_runtime_denies_expired_object_reads() {
        let runtime = InMemoryMeshRuntime::new();
        assert!(runtime.put_object(
            "mesh/object/expired",
            "AQID",
            MeshObjectPolicy {
                expires_at_ms: Some(1),
                suppress_previews: None,
            },
        ));

        assert_eq!(
            runtime.get_object_result("mesh/object/expired"),
            MeshObjectResult::unavailable(MeshObjectReadReason::Expired, Some(1))
        );
        assert_eq!(runtime.get_object("mesh/object/expired"), None);
    }

    #[test]
    fn in_memory_runtime_reports_missing_object() {
        let runtime = InMemoryMeshRuntime::new();

        assert_eq!(
            runtime.get_object_result("mesh/object/missing"),
            MeshObjectResult::unavailable(MeshObjectReadReason::NotFound, None)
        );
        assert_eq!(runtime.get_object("mesh/object/missing"), None);
    }

    #[test]
    fn test_respond_with_unknown_request_id_returns_false() {
        let runtime = InMemoryMeshRuntime::new();
        assert!(
            !runtime.respond("mesh-req-nonexistent-99", "AQID"),
            "respond with unknown request_id must return false"
        );
    }

    #[test]
    fn in_memory_runtime_supports_private_object_reads_and_repairs() {
        let runtime = InMemoryMeshRuntime::new();
        let reference = runtime.put_private_object(
            "AQID",
            MeshObjectPolicy {
                expires_at_ms: Some(u64::MAX),
                suppress_previews: Some(true),
            },
        );

        assert_eq!(
            runtime.get_private_object_result(&reference.handle, &reference.capability),
            MeshObjectResult::found("AQID".into(), Some(u64::MAX))
        );

        let request_id = runtime.request_private(&reference.handle, &reference.capability, "");
        let events = runtime.drain_events();
        assert_eq!(events.len(), 1);
        assert_eq!(events[0].event, "meshReply");
        assert_eq!(
            events[0].payload_json,
            format!(
                r#"{{"requestId":"{request_id}","dataBase64":"AQID","expiresAtMs":18446744073709551615}}"#
            )
        );
    }

    #[test]
    fn test_two_in_memory_runtimes_generate_non_colliding_request_ids() {
        let runtime_a = InMemoryMeshRuntime::new();
        let runtime_b = InMemoryMeshRuntime::new();

        let id_a = runtime_a.request("mesh/object/1", "");
        let id_b = runtime_b.request("mesh/object/1", "");

        assert_ne!(id_a, id_b, "different runtimes at same counter must differ");

        let nonce_a = id_a
            .split('-')
            .nth(2)
            .expect("nonce segment must be present");
        let nonce_b = id_b
            .split('-')
            .nth(2)
            .expect("nonce segment must be present");
        assert_ne!(nonce_a, nonce_b, "session nonces must differ");
    }

    #[test]
    fn in_memory_runtime_denies_guessed_private_access() {
        let runtime = InMemoryMeshRuntime::new();
        let reference = runtime.put_private_object("AQID", MeshObjectPolicy::default());

        assert_eq!(
            runtime.get_private_object_result(&reference.handle, "wrong-capability"),
            MeshObjectResult::unavailable(MeshObjectReadReason::PolicyDenied, None)
        );
        assert_eq!(
            runtime.get_private_object_result("guessed-handle", "guessed-capability"),
            MeshObjectResult::unavailable(MeshObjectReadReason::PolicyDenied, None)
        );
    }

    #[test]
    fn in_memory_runtime_expires_private_capabilities() {
        let runtime = InMemoryMeshRuntime::new();
        let reference = runtime.put_private_object(
            "AQID",
            MeshObjectPolicy {
                expires_at_ms: Some(1),
                suppress_previews: None,
            },
        );

        assert_eq!(
            runtime.get_private_object_result(&reference.handle, &reference.capability),
            MeshObjectResult::unavailable(MeshObjectReadReason::Expired, Some(1))
        );

        let request_id = runtime.request_private(&reference.handle, &reference.capability, "");
        let events = runtime.drain_events();
        assert_eq!(
            events[0].payload_json,
            format!(
                r#"{{"requestId":"{request_id}","dataBase64":null,"reason":"expired","expiresAtMs":1}}"#
            )
        );
    }
}