-
-
Save austin-searchpilot/d7ecdd732fa6240364c09d5e2919d374 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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