use std::{collections::HashMap, error::Error}; use html::{Rewriter, VisitorExternalTool, rule::RewriteRule}; use js_sys::{Array, Function, Object, Uint8Array}; use oxc::allocator::Allocator; use wasm_bindgen::{JsCast, JsValue, prelude::wasm_bindgen}; use crate::{ error::{Result, RewriterError}, get_obj, set_obj, }; #[wasm_bindgen(typescript_custom_section)] const REWRITER_OUTPUT: &'static str = r#" export type HtmlRewriterOutput = { html: Uint8Array, }; "#; #[wasm_bindgen] extern "C" { #[wasm_bindgen(typescript_type = "HtmlRewriterOutput")] pub type HtmlRewriterOutput; } pub type HtmlRewriter = Rewriter<(Object, Object, Function, Function, Function, Function)>; fn build_rules( rules: Vec, ) -> Result>> { rules .into_iter() .map(|x| { let entries = Object::entries(&x) .to_vec() .into_iter() .map(|x| { let arr = x .dyn_into::() .map_err(|_| RewriterError::not_arr("Object.entries"))? .to_vec(); Ok(( arr[0] .as_string() .ok_or(RewriterError::not_str("Object.entries key"))?, arr[1].clone(), )) }) .collect::>>()?; let mut func = None; let mut attrs = HashMap::new(); for (k, v) in entries { if k == "fn" { func = Some( v.dyn_into::() .map_err(|_| RewriterError::not_fn("htmlRules fn"))?, ); } else if let Ok(v) = v.clone().dyn_into::() { let els = v .to_vec() .into_iter() .map(|x| { x.as_string() .ok_or(RewriterError::not_str("htmlRules value array value")) }) .collect::>()?; attrs.insert(k, Some(els)); } else { let el = v .as_string() .ok_or(RewriterError::not_str("htmlRules value"))?; if el == "*" { attrs.insert(k, None); } } } Ok(RewriteRule { attrs, func: Box::new( move |alloc, val, (meta, cookie, _, _, _, _): &( Object, Object, Function, Function, Function, Function, )| { func.as_ref() .map(|x| { let ret = x .call3(&JsValue::NULL, &val.into(), meta, cookie) .map_err(RewriterError::from) .map_err(Box::new)?; Ok(ret.as_string().map(|x| alloc.alloc_str(&x))) }) .transpose() .map(Option::flatten) }, ), }) }) .collect::>() } #[wasm_bindgen(inline_js = r#" export function setMeta(meta, val) { meta.base = new URL(val, meta.origin); } export function base64(bytes) { const binString = Array.from(bytes, (byte) => String.fromCodePoint(byte) ).join(""); return btoa(binString); } export function rewriteJsAttr(fn, attr, code, meta) { return fn(code, `(inline ${attr} on element)`, meta); } export function rewriteJsInline(fn, code, module, meta) { return fn(code, `(inline script element)`, meta, module); } export function rewriteCss(fn, code, meta) { return fn(code, meta); } export function rewriteHttpEquiv(fn, content, meta) { const contentArray = content.split("url="); if (contentArray[1]) contentArray[1] = fn(contentArray[1].trim(), meta); return contentArray.join("url="); } export function getInjectCode(func, cookie, found_head) { return func(cookie, found_head); } export function log(val) { console.log("aaaaa", val) } "#)] extern "C" { #[wasm_bindgen(js_name = "setMeta")] fn __external_tool_set_meta(meta: &Object, val: &str); #[wasm_bindgen(js_name = "base64")] fn __external_tool_base64(bytes: Uint8Array) -> String; #[wasm_bindgen(js_name = "rewriteJsAttr")] fn __external_tool_rewrite_js_attr( func: &Function, attr: &str, code: &str, meta: &Object, ) -> String; #[wasm_bindgen(js_name = "rewriteJsInline")] fn __external_tool_rewrite_js_inline( func: &Function, code: &str, module: bool, meta: &Object, ) -> String; #[wasm_bindgen(js_name = "rewriteCss")] fn __external_tool_rewrite_css(func: &Function, code: &str, meta: &Object) -> String; #[wasm_bindgen(js_name = "rewriteHttpEquiv")] fn __external_tool_rewrite_http_equiv(func: &Function, content: &str, meta: &Object) -> String; #[wasm_bindgen(js_name = "getInjectCode")] fn __external_tool_get_inject_code( func: &Function, cookie: &Object, found_head: bool, ) -> String; #[wasm_bindgen(js_name = "log")] fn __external_tool_log(val: &str); } fn external_tool<'alloc, 'data>( alloc: &'alloc Allocator, tool: VisitorExternalTool<'data>, (meta, cookie, rewrite_js, rewrite_css, rewrite_url, get_inject): &'data ( Object, Object, Function, Function, Function, Function, ), ) -> std::result::Result, Box> { match tool { VisitorExternalTool::SetMetaBase(val) => { __external_tool_set_meta(meta, val); Ok(None) } VisitorExternalTool::Base64(val) => Ok(Some( alloc.alloc_str(&__external_tool_base64(Uint8Array::from(val.as_bytes()))), )), VisitorExternalTool::RewriteJsAttr { attr, code } => Ok(Some(alloc.alloc_str( &__external_tool_rewrite_js_attr(rewrite_js, attr, code, meta), ))), VisitorExternalTool::RewriteInlineScript { code, module } => Ok(Some(alloc.alloc_str( &__external_tool_rewrite_js_inline(rewrite_js, code, module, meta), ))), VisitorExternalTool::RewriteCss(css) => Ok(Some( alloc.alloc_str(&__external_tool_rewrite_css(rewrite_css, css, meta)), )), VisitorExternalTool::RewriteHttpEquivContent(content) => Ok(Some(alloc.alloc_str( &__external_tool_rewrite_http_equiv(rewrite_url, content, meta), ))), VisitorExternalTool::Log(log) => { __external_tool_log(log); Ok(None) } VisitorExternalTool::GetScriptText { found_head } => Ok(Some(alloc.alloc_str( &__external_tool_get_inject_code(get_inject, cookie, found_head), ))), } } pub fn create_html(scramjet: &Object) -> Result { let shared = get_obj(scramjet, "shared")?; let rewrite = get_obj(&shared, "rewrite")?; let html_rules = get_obj(&rewrite, "htmlRules")? .dyn_into::() .map_err(|_| RewriterError::not_arr("htmlRules"))?; let html_rules = html_rules.to_vec().into_iter().map(Object::from).collect(); let rules = build_rules(html_rules)?; Ok(HtmlRewriter::new(rules, Box::new(external_tool))?) } pub fn get_html_params(scramjet: &Object) -> Result<(Function, Function, Function, Function)> { let shared = get_obj(scramjet, "shared")?; let rewrite = get_obj(&shared, "rewrite")?; let js = get_obj(&rewrite, "rewriteJs")? .dyn_into::() .map_err(|_| RewriterError::not_fn("rewriteJs"))?; let css = get_obj(&rewrite, "rewriteCss")? .dyn_into::() .map_err(|_| RewriterError::not_fn("rewriteCss"))?; let url = get_obj(&rewrite, "rewriteUrl")? .dyn_into::() .map_err(|_| RewriterError::not_fn("rewriteUrl"))?; let inject_code = get_obj(&rewrite, "getHtmlInjectCode")? .dyn_into::() .map_err(|_| RewriterError::not_fn("getHtmlInjectCode"))?; Ok((js, css, url, inject_code)) } pub fn create_html_output(out: &[u8]) -> Result { let obj = Object::new(); set_obj(&obj, "html", &Uint8Array::from(out).into())?; Ok(HtmlRewriterOutput::from(JsValue::from(obj))) }