claw/src/browser/ws_protocol.rs

use serde_json::{json, Value};

use crate::pipe::{Action, PipeError};

const CALLBACK_DELIMITER: &str = "@_@";
const CALLBACK_PREFIX: &str = "sgclaw_cb_";
const JS_AREA_HIDE: &str = "hide";

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CallbackCorrelation {
    pub request_id: String,
    pub callback_name: String,
    pub source_url: String,
    pub target_url: String,
    pub action_url: String,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct EncodedWsRequest {
    pub payload: String,
    pub callback: Option<CallbackCorrelation>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DecodedCallback {
    pub source_url: String,
    pub target_url: String,
    pub callback_name: String,
    pub action_url: String,
    pub response_text: String,
}

pub fn encode_v1_action(
    action: &Action,
    params: &Value,
    request_url: &str,
    request_id: Option<&str>,
) -> Result<EncodedWsRequest, PipeError> {
    match action {
        Action::Navigate => encode_navigate(params, request_url, request_id),
        Action::Click => encode_click(params, request_url),
        Action::Type => encode_type(params, request_url),
        Action::GetText => encode_get_text(params, request_url, request_id),
        Action::Eval => encode_eval(params, request_url, request_id),
        _ => Err(PipeError::Protocol(format!(
            "unsupported browser ws action: {}",
            action.as_str()
        ))),
    }
}

pub fn decode_callback_frame(frame: &str) -> Result<DecodedCallback, PipeError> {
    let payload: Value = serde_json::from_str(frame)?;
    let array = payload.as_array().ok_or_else(|| {
        PipeError::Protocol("callback frame must be a JSON array".to_string())
    })?;
    if array.len() != 3 {
        return Err(PipeError::Protocol(
            "callback frame must contain [requesturl, function, payload]".to_string(),
        ));
    }

    let function_name = array[1].as_str().ok_or_else(|| {
        PipeError::Protocol("callback frame function name must be a string".to_string())
    })?;
    if function_name != "callBackJsToCpp" {
        return Err(PipeError::Protocol(
            "callback frame must target callBackJsToCpp".to_string(),
        ));
    }

    let param = array[2].as_str().ok_or_else(|| {
        PipeError::Protocol("callback payload must be a string".to_string())
    })?;
    let mut parts = param.splitn(5, CALLBACK_DELIMITER);
    let source_url = parts.next().unwrap_or_default();
    let target_url = parts.next().unwrap_or_default();
    let callback_name = parts.next().unwrap_or_default();
    let action_url = parts.next().unwrap_or_default();
    let response_text = parts.next().unwrap_or_default();

    if source_url.is_empty()
        || target_url.is_empty()
        || callback_name.is_empty()
        || action_url.is_empty()
        || response_text.is_empty() && !param.ends_with(CALLBACK_DELIMITER)
    {
        return Err(PipeError::Protocol(
            "malformed callback payload".to_string(),
        ));
    }

    Ok(DecodedCallback {
        source_url: source_url.to_string(),
        target_url: target_url.to_string(),
        callback_name: callback_name.to_string(),
        action_url: action_url.to_string(),
        response_text: response_text.to_string(),
    })
}

fn encode_navigate(
    params: &Value,
    request_url: &str,
    request_id: Option<&str>,
) -> Result<EncodedWsRequest, PipeError> {
    let url = required_string(params, "url")?;
    let callback = callback_metadata(
        request_id,
        request_url,
        &url,
        "sgHideBrowserCallAfterLoaded",
    )?;
    let callback_call = format!(
        "callBackJsToCpp(\"{request_url}@_@{url}@_@{callback_name}@_@sgHideBrowserCallAfterLoaded@_@\")",
        callback_name = callback.callback_name,
    );
    Ok(EncodedWsRequest {
        payload: serde_json::to_string(&json!([
            request_url,
            "sgHideBrowserCallAfterLoaded",
            url,
            callback_call,
        ]))?,
        callback: Some(callback),
    })
}

fn encode_click(params: &Value, request_url: &str) -> Result<EncodedWsRequest, PipeError> {
    let target_url = target_url(params, request_url)?;
    let selector = required_string(params, "selector")?;
    let script = format!(
        "(function(){{const el=document.querySelector({selector:?});if(!el){{throw new Error(\"selector not found: {selector}\");}}el.click();}})();"
    );
    encode_js_in_area(request_url, &target_url, &script, None)
}

fn encode_type(params: &Value, request_url: &str) -> Result<EncodedWsRequest, PipeError> {
    let target_url = target_url(params, request_url)?;
    let selector = required_string(params, "selector")?;
    let text = required_string(params, "text")?;
    let script = format!(
        "(function(){{const el=document.querySelector({selector:?});if(!el){{throw new Error(\"selector not found: {selector}\");}}el.value={text:?};el.dispatchEvent(new Event(\"input\",{{bubbles:true}}));el.dispatchEvent(new Event(\"change\",{{bubbles:true}}));}})();"
    );
    encode_js_in_area(request_url, &target_url, &script, None)
}

fn encode_get_text(
    params: &Value,
    request_url: &str,
    request_id: Option<&str>,
) -> Result<EncodedWsRequest, PipeError> {
    let target_url = target_url(params, request_url)?;
    let selector = required_string(params, "selector")?;
    let callback = callback_metadata(
        request_id,
        request_url,
        &target_url,
        "sgBrowserExcuteJsCodeByArea",
    )?;
    let script = format!(
        "(function(){{const el=document.querySelector({selector:?});if(!el){{throw new Error(\"selector not found: {selector}\");}}const text=el.innerText ?? el.textContent ?? \"\";callBackJsToCpp(\"{request_url}@_@{target_url}@_@{callback_name}@_@sgBrowserExcuteJsCodeByArea@_@\"+String(text));}})();",
        callback_name = callback.callback_name
    );
    encode_js_in_area(request_url, &target_url, &script, Some(callback))
}

fn encode_eval(
    params: &Value,
    request_url: &str,
    request_id: Option<&str>,
) -> Result<EncodedWsRequest, PipeError> {
    let target_url = target_url(params, request_url)?;
    let source_script = required_string(params, "script")?;
    let callback = callback_metadata(
        request_id,
        request_url,
        &target_url,
        "sgBrowserExcuteJsCodeByArea",
    )?;
    let script = format!(
        "(function(){{const result=(function(){{{source_script}}})();callBackJsToCpp(\"{request_url}@_@{target_url}@_@{callback_name}@_@sgBrowserExcuteJsCodeByArea@_@\"+String(result));}})();",
        callback_name = callback.callback_name
    );
    encode_js_in_area(request_url, &target_url, &script, Some(callback))
}

fn encode_js_in_area(
    request_url: &str,
    target_url: &str,
    script: &str,
    callback: Option<CallbackCorrelation>,
) -> Result<EncodedWsRequest, PipeError> {
    Ok(EncodedWsRequest {
        payload: serde_json::to_string(&json!([
            request_url,
            "sgBrowserExcuteJsCodeByArea",
            target_url,
            script,
            JS_AREA_HIDE,
        ]))?,
        callback,
    })
}

fn callback_metadata(
    request_id: Option<&str>,
    request_url: &str,
    target_url: &str,
    action_url: &str,
) -> Result<CallbackCorrelation, PipeError> {
    let request_id = request_id
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .ok_or_else(|| PipeError::Protocol("request_id is required".to_string()))?;
    Ok(CallbackCorrelation {
        request_id: request_id.to_string(),
        callback_name: format!("{CALLBACK_PREFIX}{request_id}"),
        source_url: request_url.to_string(),
        target_url: target_url.to_string(),
        action_url: action_url.to_string(),
    })
}

fn target_url(params: &Value, request_url: &str) -> Result<String, PipeError> {
    Ok(optional_string(params, "target_url")
        .filter(|value| !value.trim().is_empty())
        .unwrap_or_else(|| request_url.to_string()))
}

fn required_string(params: &Value, key: &str) -> Result<String, PipeError> {
    optional_string(params, key)
        .filter(|value| !value.trim().is_empty())
        .ok_or_else(|| PipeError::Protocol(format!("{key} is required")))
}

fn optional_string(params: &Value, key: &str) -> Option<String> {
    params.get(key)?.as_str().map(ToString::to_string)
}

#[cfg(test)]
mod tests {
    use super::{decode_callback_frame, encode_v1_action};
    use crate::pipe::Action;
    use serde_json::{json, Value};

    #[test]
    fn get_text_callback_uses_documented_browser_opcode() {
        let request = encode_v1_action(
            &Action::GetText,
            &json!({
                "target_url": "https://www.zhihu.com/hot",
                "selector": "#content"
            }),
            "https://www.zhihu.com/hot",
            Some("req42"),
        )
        .unwrap();

        let payload: Value = serde_json::from_str(&request.payload).unwrap();
        assert_eq!(payload[1], json!("sgBrowserExcuteJsCodeByArea"));
        assert_eq!(payload[4], json!("hide"));
        assert_eq!(
            request.callback.unwrap().action_url,
            "sgBrowserExcuteJsCodeByArea"
        );
        assert!(payload[3].as_str().unwrap().contains(
            "callBackJsToCpp(\"https://www.zhihu.com/hot@_@https://www.zhihu.com/hot@_@sgclaw_cb_req42@_@sgBrowserExcuteJsCodeByArea@_@\"+String(text))"
        ));
    }

    #[test]
    fn eval_callback_uses_documented_browser_opcode() {
        let request = encode_v1_action(
            &Action::Eval,
            &json!({
                "target_url": "https://www.zhihu.com/hot",
                "script": "2 + 2"
            }),
            "https://www.zhihu.com/hot",
            Some("req-eval"),
        )
        .unwrap();

        let payload: Value = serde_json::from_str(&request.payload).unwrap();
        assert_eq!(payload[1], json!("sgBrowserExcuteJsCodeByArea"));
        assert_eq!(
            request.callback.unwrap().action_url,
            "sgBrowserExcuteJsCodeByArea"
        );
        assert!(payload[3].as_str().unwrap().contains(
            "callBackJsToCpp(\"https://www.zhihu.com/hot@_@https://www.zhihu.com/hot@_@sgclaw_cb_req-eval@_@sgBrowserExcuteJsCodeByArea@_@\"+String(result))"
        ));
    }

    #[test]
    fn decodes_documented_callback_payload() {
        let callback = decode_callback_frame(
            r#"["https://www.zhihu.com/hot","callBackJsToCpp","https://www.zhihu.com/hot@_@https://www.zhihu.com/hot@_@sgclaw_cb_req42@_@sgBrowserExcuteJsCodeByArea@_@天气"]"#,
        )
        .unwrap();

        assert_eq!(callback.action_url, "sgBrowserExcuteJsCodeByArea");
        assert_eq!(callback.response_text, "天气");
    }
}