Skip to content

Instantly share code, notes, and snippets.

@austin-searchpilot
Created June 10, 2021 14:53
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save austin-searchpilot/d7ecdd732fa6240364c09d5e2919d374 to your computer and use it in GitHub Desktop.
Save austin-searchpilot/d7ecdd732fa6240364c09d5e2919d374 to your computer and use it in GitHub Desktop.
extern crate cfg_if;
extern crate wasm_bindgen;
mod utils;
use std::pin::Pin;
use cfg_if::cfg_if;
// use futures::StreamExt;
use lol_html::html_content::ContentType;
use lol_html::OutputSink;
use lol_html::{element, rewrite_str, HtmlRewriter, RewriteStrSettings, Settings};
use wasm_bindgen::prelude::*;
use wasm_bindgen::prelude::*;
use wasm_bindgen::JsCast;
use wasm_bindgen_futures::{spawn_local, JsFuture};
use web_sys::{Request, RequestInit, Response};
// use web_sys::Request;
use futures::{SinkExt, StreamExt};
// use futures_sink::Sink;
use js_sys::Uint8Array;
use wasm_streams::transform::TransformStream;
use wasm_streams::writable::{WritableStream, WritableStreamDefaultWriter};
use wasm_streams::ReadableStream;
use web_sys::console;
// Returns global execution context of a service worker
fn worker_global_scope() -> Option<web_sys::ServiceWorkerGlobalScope> {
js_sys::global()
.dyn_into::<web_sys::ServiceWorkerGlobalScope>()
.ok()
}
#[wasm_bindgen]
pub async fn modify_html(request: web_sys::Request) -> Response {
let global = worker_global_scope().unwrap();
let resp_value = JsFuture::from(global.fetch_with_request(&request))
.await
.unwrap();
assert!(resp_value.is_instance_of::<Response>());
let response: Response = resp_value.dyn_into().unwrap();
// Get the response's body as a JS ReadableStream
let raw_body = response.body().unwrap_throw();
let body = ReadableStream::from_raw(raw_body.dyn_into().unwrap_throw());
// Convert the JS ReadableStream to a Rust stream
let mut stream = body.into_stream();
let output_stream_raw = wasm_streams::transform::sys::TransformStream::new();
let output_stream = wasm_streams::transform::TransformStream::from_raw(output_stream_raw);
let new_body_reader = output_stream.readable();
let new_body_reader_raw = new_body_reader.into_raw();
let new_body_reader_web = new_body_reader_raw.unchecked_into();
let new_response = Response::new_with_opt_readable_stream(Some(&new_body_reader_web)).unwrap();
spawn_local(async move {
// initialise the body writer
let mut new_body_writer = output_stream.writable();
struct StreamWriter {
w: Pin<Box<WritableStream>>,
}
impl OutputSink for StreamWriter {
fn handle_chunk(&mut self, chunk: &[u8]) {
let to_write = Uint8Array::from(chunk);
let new_body_writer_writer = self.w.get_writer();
let sink = new_body_writer_writer.into_sink();
spawn_local(async move {
sink.send(to_write.into()).await;
});
}
}
let output_sink = StreamWriter {
w: Box::pin(new_body_writer),
};
let mut rewriter = HtmlRewriter::try_new(
Settings {
element_content_handlers: vec![
// Rewrite title tag
element!("title", |el| {
el.set_inner_content("Hello from CF rust!", ContentType::Text);
Ok(())
}),
],
..Settings::default()
},
output_sink,
)
.unwrap();
// Consume the stream, logging each individual chunk
while let Some(Ok(chunk)) = stream.next().await {
console::log_1(&chunk);
let chunk_array: Uint8Array = chunk.dyn_into().unwrap();
let rust_chunk = chunk_array.to_vec();
rewriter.write(&rust_chunk).unwrap();
}
});
return new_response;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment