workflow-nw 0.19.0

use crate::ipc::imports::*;
use crate::ipc::method::*;
use crate::ipc::notification::*;
use crate::ipc::target::*;

/// Identifier type used to correlate IPC requests with their responses.
pub type IpcId = Id64;

struct Pending<F> {
    _timestamp: Instant,
    callback: F,
}
impl<F> Pending<F> {
    fn new(callback: F) -> Self {
        Self {
            _timestamp: Instant::now(),
            callback,
        }
    }
}

type PendingMap<Id, F> = Arc<Mutex<AHashMap<Id, Pending<F>>>>;

/// Callback type invoked to deliver a Borsh-encoded response (or error) back
/// to the originator of a pending IPC request.
pub type BorshResponseFn = Arc<
    Box<dyn Fn(Vec<u8>, ResponseResult<Vec<u8>>, Option<&Duration>) -> Result<()> + Sync + Send>,
>;

struct Inner<Ops>
where
    Ops: OpsT,
{
    target: IpcTarget,
    identifier: String,
    handler: Mutex<Option<Arc<dyn AsCallback>>>,
    methods: Mutex<AHashMap<Ops, Arc<dyn MethodTrait>>>,
    notifications: Mutex<AHashMap<Ops, Arc<dyn NotificationTrait>>>,
}

/// IPC endpoint that handles inbound messages for a target context and
/// maintains the registered method and notification handlers, keyed by the
/// operation type `Ops`.
pub struct Ipc<Ops>
where
    Ops: OpsT,
{
    inner: Arc<Inner<Ops>>,
    _ops: PhantomData<Ops>,
}

unsafe impl<Ops> Send for Ipc<Ops> where Ops: OpsT {}
unsafe impl<Ops> Sync for Ipc<Ops> where Ops: OpsT {}

impl<Ops> Drop for Ipc<Ops>
where
    Ops: OpsT,
{
    fn drop(&mut self) {
        self.unregister_handler().ok();
    }
}

impl<Ops> Ipc<Ops>
where
    Ops: OpsT,
{
    /// Creates an IPC binding bound to the global context, registering it as
    /// the global IPC handler source. Panics if a handler source is already
    /// registered.
    pub fn try_new_global_binding<Ident>(identifier: Ident) -> Result<Arc<Self>>
    where
        Ident: ToString,
    {
        let target = IpcTarget::new(global::global().as_ref());
        let ipc = Self::try_new_binding(&target, identifier)?;

        let ipc_handler_source_ptr = &raw mut IPC_HANDLER_SOURCE;

        unsafe {
            if (*ipc_handler_source_ptr).is_some() {
                panic!("global ipc handler already registered");
            }
            (*ipc_handler_source_ptr).replace(target);
        }

        Ok(ipc)
    }

    /// Creates an IPC binding bound to the given window's context, registering
    /// it as the global IPC handler source. Panics if a handler source is
    /// already registered.
    pub fn try_new_window_binding<Ident>(
        window: &Arc<Window>,
        identifier: Ident,
    ) -> Result<Arc<Self>>
    where
        Ident: ToString,
    {
        let window = window.window();
        let target = IpcTarget::new(window.as_ref());
        let ipc = Self::try_new_binding(&target, identifier)?;

        let ipc_handler_source_ptr = &raw mut IPC_HANDLER_SOURCE;

        unsafe {
            if (*ipc_handler_source_ptr).is_some() {
                panic!("global ipc handler already registered");
            }
            (*ipc_handler_source_ptr).replace(target);
        }

        Ok(ipc)
    }

    fn try_new_binding<Ident>(target: &IpcTarget, identifier: Ident) -> Result<Arc<Self>>
    where
        Ident: ToString,
    {
        let ipc = Arc::new(Ipc {
            inner: Arc::new(Inner {
                target: target.clone(),
                identifier: identifier.to_string(),
                handler: Mutex::new(None),
                methods: Mutex::new(AHashMap::default()),
                notifications: Mutex::new(AHashMap::default()),
            }),
            _ops: PhantomData,
        });

        ipc.register_handler()?;

        Ok(ipc)
    }

    fn register_handler(self: &Arc<Self>) -> Result<()> {
        let this = self.clone();
        let handler = Arc::new(callback!(move |message: ArrayBuffer, source: JsValue| {
            let this = this.clone();

            let message = Uint8Array::new(&message);
            let vec = message.to_vec();

            let source = if source == JsValue::NULL {
                None
            } else {
                Some(IpcTarget::new(source.as_ref()))
            };

            spawn(async move {
                match BorshMessage::<IpcId>::try_from(&vec) {
                    Ok(message) => {
                        if let Err(err) = this.handle_message(message, source).await {
                            log_error!("IPC: handler error: {:?}", err);
                        }
                    }
                    Err(err) => {
                        log_error!("Failed to deserialize ipc message: {:?}", err);
                    }
                }
            })
        }));

        js_sys::Reflect::set(
            self.inner.target.as_ref(),
            &JsValue::from_str("ipc_handler"),
            handler.get_fn(),
        )?;
        js_sys::Reflect::set(
            self.inner.target.as_ref(),
            &JsValue::from_str("ipc_identifier"),
            &JsValue::from(&self.inner.identifier),
        )?;

        self.inner.handler.lock().unwrap().replace(handler);

        Ok(())
    }

    fn unregister_handler(&self) -> Result<()> {
        if let Some(_handler) = self.inner.handler.lock().unwrap().take() {
            let object = Object::from(self.inner.target.as_ref().clone());
            js_sys::Reflect::delete_property(&object, &JsValue::from_str("ipc_handler"))?;
            js_sys::Reflect::delete_property(&object, &JsValue::from_str("ipc_identifier"))?;
        }

        Ok(())
    }

    /// Processes an incoming Borsh-encoded IPC message, dispatching it to the
    /// appropriate registered method or notification handler, or resolving a
    /// pending outbound call when the message is a response.
    pub async fn handle_message(
        &self,
        message: BorshMessage<'_, IpcId>,
        source: Option<IpcTarget>,
    ) -> Result<()> {
        let BorshMessage::<IpcId> { header, payload } = message;
        let BorshHeader::<IpcId> { op, id, kind } = header;
        match kind {
            MessageKind::Request => {
                let source = source.unwrap_or_else(|| {
                    panic!("ipc received a call request with no source: {:?}", op)
                });

                let op = Ops::try_from_slice(&op)?;

                let method = self.inner.methods.lock().unwrap().get(&op).cloned();
                if let Some(method) = method {
                    let result = method.call_with_borsh(payload).await;
                    let buffer = borsh::to_vec(&result)?;
                    source.call_ipc(
                        to_msg::<Ops, IpcId>(BorshHeader::response(id, op), &buffer)?.as_ref(),
                        None,
                    )?;
                } else {
                    log_error!("ipc method handler not found: {:?}", op);
                    let resp: ResponseResult<()> = Err(ResponseError::NotFound);
                    let buffer = borsh::to_vec(&resp)?;
                    source.call_ipc(
                        to_msg::<Ops, IpcId>(BorshHeader::response(id, op), &buffer)?.as_ref(),
                        None,
                    )?;
                }
            }
            MessageKind::Notification => {
                let op = Ops::try_from_slice(&op)?;

                let notification = self.inner.notifications.lock().unwrap().get(&op).cloned();

                if let Some(notification) = notification {
                    match notification.call_with_borsh(payload).await {
                        Ok(_resp) => {}
                        Err(err) => {
                            log_error!("ipc notification error: {:?}", err);
                        }
                    }
                } else {
                    log_error!("ipc notification handler not found: {:?}", op);
                }
            }
            MessageKind::Response => {
                let id = id.expect("ipc missing success response id");
                // let id = Id64::from(id);
                let mut pending = pending().lock().unwrap();
                match pending.remove(&id) {
                    Some(pending) => {
                        let resp = ResponseResult::<Vec<u8>>::try_from_slice(payload)?;
                        (pending.callback)(op, resp, None)?;
                    }
                    _ => {
                        log_error!("ipc response id not found: {:?}", id);
                    }
                }
            }
        }

        Ok(())
    }

    /// Registers a request/response method handler for the given operation.
    /// Panics if a handler for the same operation has already been registered.
    pub fn method<Req, Resp>(&self, op: Ops, method: Method<Req, Resp>)
    where
        Ops: Debug + Clone,
        Req: MsgT,
        Resp: MsgT,
    {
        let method: Arc<dyn MethodTrait> = Arc::new(method);
        if self
            .inner
            .methods
            .lock()
            .unwrap()
            .insert(op.clone(), method)
            .is_some()
        {
            panic!("RPC method {op:?} is declared multiple times")
        }
    }

    /// Registers a notification handler for the given operation. Panics if a
    /// handler for the same operation has already been registered.
    pub fn notification<Msg>(&self, op: Ops, method: Notification<Msg>)
    where
        Ops: Debug + Clone,
        Msg: MsgT,
    {
        let method: Arc<dyn NotificationTrait> = Arc::new(method);
        if self
            .inner
            .notifications
            .lock()
            .unwrap()
            .insert(op.clone(), method)
            .is_some()
        {
            panic!("RPC notification {op:?} is declared multiple times")
        }
    }
}

trait IpcHandler {
    fn call_ipc(&self, data: &JsValue, source: Option<&IpcTarget>) -> Result<()>;
}

impl IpcHandler for IpcTarget {
    fn call_ipc(&self, data: &JsValue, source: Option<&IpcTarget>) -> Result<()> {
        let target_fn = js_sys::Reflect::get(self.as_ref(), &JsValue::from_str("ipc_handler"))?;

        let target_fn = target_fn.unchecked_into::<js_sys::Function>();

        if let Some(source) = source {
            target_fn.call2(
                &JsValue::UNDEFINED,
                &JsValue::from(data),
                &JsValue::from(source.as_ref()),
            )?;
        } else {
            target_fn.call2(&JsValue::UNDEFINED, &JsValue::from(data), &JsValue::NULL)?;
        }

        Ok(())
    }
}

static mut PENDING: Option<PendingMap<IpcId, BorshResponseFn>> = None; //PendingMap::default();
fn pending() -> &'static mut PendingMap<IpcId, BorshResponseFn> {
    let pending_ptr = &raw mut PENDING;
    unsafe {
        if (*pending_ptr).is_none() {
            PENDING = Some(PendingMap::default());
        }
        (*pending_ptr).as_mut().unwrap()
    }
}

static mut IPC_HANDLER_SOURCE: Option<IpcTarget> = None;

/// Trait implemented by types that can act as an IPC peer, providing the
/// ability to send notifications and issue request/response calls to a
/// target context (a window or the global object).
#[async_trait]
pub trait IpcDispatch {
    /// Returns the [`IpcTarget`] that messages dispatched through this peer
    /// should be delivered to.
    fn as_target(&self) -> IpcTarget;

    /// Sends a one-way notification carrying `op` and a Borsh-serialized
    /// `payload` to the target context, without awaiting a response.
    async fn notify<Ops, Msg>(&self, op: Ops, payload: Msg) -> Result<()>
    where
        Ops: OpsT,
        Msg: BorshSerialize + Send + Sync + 'static,
    {
        let payload = borsh::to_vec(&payload).map_err(|_| Error::BorshSerialize)?;
        self.as_target().call_ipc(
            to_msg::<Ops, IpcId>(BorshHeader::notification::<Ops>(op), &payload)?.as_ref(),
            None,
        )?;
        Ok(())
    }

    /// Issues a request/response call carrying `op` and `req`, awaiting and
    /// deserializing the response. Uses the registered local IPC object as the
    /// reply source.
    async fn call<Ops, Req, Resp>(&self, op: Ops, req: Req) -> Result<Resp>
    where
        Ops: OpsT,
        Req: MsgT,
        Resp: MsgT,
    {
        let ipc_handler_source_ptr = &raw const IPC_HANDLER_SOURCE;
        let source = unsafe {
            (*ipc_handler_source_ptr)
                .as_ref()
                .cloned()
                .expect("missing ipc handler source (please register a local IPC object)")
        };
        self.call_with_source(op, req, &source).await
    }

    /// Like [`call`](Self::call), but routes the response to an explicit
    /// `source` [`IpcTarget`] instead of the registered local IPC object.
    async fn call_with_source<Ops, Req, Resp>(
        &self,
        op: Ops,
        req: Req,
        source: &IpcTarget,
    ) -> Result<Resp>
    where
        Ops: OpsT,
        Req: MsgT,
        Resp: MsgT,
    {
        let payload = borsh::to_vec(&req).map_err(|_| Error::BorshSerialize)?;

        let id = Id64::generate();
        let (sender, receiver) = oneshot();

        {
            let mut pending = pending().lock().unwrap();
            pending.insert(
                id.clone(),
                Pending::new(Arc::new(Box::new(move |op, result, _duration| {
                    sender.try_send((op, result.map(|data| data.to_vec())))?;
                    Ok(())
                }))),
            );
        }

        self.as_target().call_ipc(
            to_msg::<Ops, IpcId>(BorshHeader::request::<Ops>(Some(id), op.clone()), &payload)?
                .as_ref(),
            Some(source),
        )?;

        let (op_, data) = receiver.recv().await?;

        let op_ = Ops::try_from_slice(&op_)?;
        if op != op_ {
            return Err(Error::Custom(format!(
                "ipc op mismatch: expected {:?}, got {:?}",
                op, op_
            )));
        }

        let resp = ResponseResult::<Resp>::try_from_slice(data?.as_ref())
            .map_err(|e| Error::BorshDeserialize(e.to_string()))?;

        Ok(resp?)
    }
}

impl IpcDispatch for IpcTarget {
    fn as_target(&self) -> IpcTarget {
        self.clone()
    }
}

impl IpcDispatch for nw_sys::Window {
    fn as_target(&self) -> IpcTarget {
        IpcTarget::new(self.window().as_ref())
    }
}

/// Locates an [`IpcTarget`] (the global context or one of the open windows)
/// whose registered IPC handler matches the given identifier, returning
/// `Ok(None)` if no such target exists.
pub async fn get_ipc_target<Ident>(identifier: Ident) -> crate::result::Result<Option<IpcTarget>>
where
    Ident: ToString,
{
    let ident: String = identifier.to_string();

    if let Some(ipc_ident) =
        Reflect::get(&global::global(), &JsValue::from("ipc_identifier"))?.as_string()
        && ipc_ident == ident
    {
        return Ok(Some(IpcTarget::new(global::global().as_ref())));
    }

    let windows = crate::window::get_all_async().await?;

    for window in windows.iter() {
        let prop =
            js_sys::Reflect::get(window.window().as_ref(), &JsValue::from("ipc_identifier"))?;
        if let Some(ipc_ident) = prop.as_string()
            && ipc_ident == ident
        {
            return Ok(Some(IpcTarget::new(window.window().as_ref())));
        }
    }
    Ok(None)
}