use super::{BrowserContext, CtxHandlerFuture, ServiceWorkerHandlerFuture};
use crate::error::Result;
use crate::protocol::EventValue;
use crate::protocol::event_registry::{EventRegistry, Handler, bridge};
use crate::protocol::event_waiter::EventWaiter;
use crate::protocol::{Download, Frame, Page, Request, ResponseObject};
use crate::server::channel_owner::ChannelOwner;
use std::future::Future;
use std::sync::{Arc, Mutex};
use tokio::sync::oneshot;
impl BrowserContext {
#[tracing::instrument(level = "debug", skip_all, fields(guid = %self.guid()))]
pub async fn on_page<F, Fut>(&self, handler: F) -> Result<()>
where
F: Fn(Page) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
let handler: Handler<Page> = Arc::new(move |page| Box::pin(handler(page)));
self.page_events.add_handler(handler);
Ok(())
}
async fn subscribe_if_idle<T>(&self, reg: &EventRegistry<T>) {
if reg.is_idle() {
_ = self.channel().update_subscription(reg.name(), true).await;
}
}
#[tracing::instrument(level = "debug", skip_all, fields(guid = %self.guid()))]
pub async fn on_download<F, Fut>(&self, handler: F) -> Result<()>
where
F: Fn(Download) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
add_forwarded(&self.forwarders.download, handler);
Ok(())
}
#[tracing::instrument(level = "debug", skip_all, fields(guid = %self.guid()))]
pub async fn on_frame_attached<F, Fut>(&self, handler: F) -> Result<()>
where
F: Fn(Frame) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
add_forwarded(&self.forwarders.frame_attached, handler);
Ok(())
}
#[tracing::instrument(level = "debug", skip_all, fields(guid = %self.guid()))]
pub async fn on_frame_detached<F, Fut>(&self, handler: F) -> Result<()>
where
F: Fn(Frame) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
add_forwarded(&self.forwarders.frame_detached, handler);
Ok(())
}
#[tracing::instrument(level = "debug", skip_all, fields(guid = %self.guid()))]
pub async fn on_frame_navigated<F, Fut>(&self, handler: F) -> Result<()>
where
F: Fn(Frame) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
add_forwarded(&self.forwarders.frame_navigated, handler);
Ok(())
}
#[tracing::instrument(level = "debug", skip_all, fields(guid = %self.guid()))]
pub async fn on_page_load<F, Fut>(&self, handler: F) -> Result<()>
where
F: Fn(Page) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
add_forwarded(&self.forwarders.page_load, handler);
Ok(())
}
#[tracing::instrument(level = "debug", skip_all, fields(guid = %self.guid()))]
pub async fn on_page_close<F, Fut>(&self, handler: F) -> Result<()>
where
F: Fn(Page) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
add_forwarded(&self.forwarders.page_close, handler);
Ok(())
}
#[tracing::instrument(level = "debug", skip_all, fields(guid = %self.guid()))]
pub async fn on_close<F, Fut>(&self, handler: F) -> Result<()>
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
let handler: Handler<()> = Arc::new(move |()| Box::pin(handler()));
self.close_events.add_handler(handler);
Ok(())
}
#[tracing::instrument(level = "debug", skip_all, fields(guid = %self.guid()))]
pub async fn on_request<F, Fut>(&self, handler: F) -> Result<()>
where
F: Fn(Request) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
let handler: Handler<Request> = Arc::new(move |request| Box::pin(handler(request)));
self.subscribe_if_idle(&self.request).await;
self.request.add_handler(handler);
Ok(())
}
#[tracing::instrument(level = "debug", skip_all, fields(guid = %self.guid()))]
pub async fn on_request_finished<F, Fut>(&self, handler: F) -> Result<()>
where
F: Fn(Request) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
let handler: Handler<Request> = Arc::new(move |request| Box::pin(handler(request)));
self.subscribe_if_idle(&self.request_finished).await;
self.request_finished.add_handler(handler);
Ok(())
}
#[tracing::instrument(level = "debug", skip_all, fields(guid = %self.guid()))]
pub async fn on_request_failed<F, Fut>(&self, handler: F) -> Result<()>
where
F: Fn(Request) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
let handler: Handler<Request> = Arc::new(move |request| Box::pin(handler(request)));
self.subscribe_if_idle(&self.request_failed).await;
self.request_failed.add_handler(handler);
Ok(())
}
#[tracing::instrument(level = "debug", skip_all, fields(guid = %self.guid()))]
pub async fn on_response<F, Fut>(&self, handler: F) -> Result<()>
where
F: Fn(ResponseObject) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
let handler: Handler<ResponseObject> =
Arc::new(move |response| Box::pin(handler(response)));
self.subscribe_if_idle(&self.response).await;
self.response.add_handler(handler);
Ok(())
}
#[tracing::instrument(level = "debug", skip_all, fields(guid = %self.guid()))]
pub async fn on_dialog<F, Fut>(&self, handler: F) -> Result<()>
where
F: Fn(crate::protocol::Dialog) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
let handler: Handler<crate::protocol::Dialog> =
Arc::new(move |dialog| Box::pin(handler(dialog)));
self.dialog.add_handler(handler);
Ok(())
}
pub async fn on_dialog_closed<F, Fut>(&self, handler: F) -> Result<()>
where
F: Fn(crate::protocol::Dialog) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
let handler: Handler<crate::protocol::Dialog> =
Arc::new(move |dialog| Box::pin(handler(dialog)));
self.subscribe_if_idle(&self.dialog_closed).await;
self.dialog_closed.add_handler(handler);
Ok(())
}
pub(crate) async fn ensure_dialog_closed_subscription(&self) {
self.subscribe_if_idle(&self.dialog_closed).await;
}
#[tracing::instrument(level = "debug", skip_all, fields(guid = %self.guid()))]
pub async fn on_console<F, Fut>(&self, handler: F) -> Result<()>
where
F: Fn(crate::protocol::ConsoleMessage) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
let handler: Handler<crate::protocol::ConsoleMessage> =
Arc::new(move |msg| Box::pin(handler(msg)));
self.subscribe_if_idle(&self.console).await;
self.console.add_handler(handler);
Ok(())
}
#[tracing::instrument(level = "debug", skip_all, fields(guid = %self.guid()))]
pub async fn on_weberror<F, Fut>(&self, handler: F) -> Result<()>
where
F: Fn(crate::protocol::WebError) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
let handler: Handler<crate::protocol::WebError> =
Arc::new(move |web_error| Box::pin(handler(web_error)));
self.weberror.add_handler(handler);
Ok(())
}
#[tracing::instrument(level = "debug", skip_all, fields(guid = %self.guid()))]
pub async fn on_serviceworker<F, Fut>(&self, handler: F) -> Result<()>
where
F: Fn(crate::protocol::Worker) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
let handler = Arc::new(
move |worker: crate::protocol::Worker| -> ServiceWorkerHandlerFuture {
Box::pin(handler(worker))
},
);
self.serviceworker_handlers.lock().unwrap().push(handler);
Ok(())
}
#[tracing::instrument(level = "debug", skip_all, fields(guid = %self.guid()))]
pub async fn expect_page(&self, timeout: Option<f64>) -> Result<EventWaiter<Page>> {
let rx = self.page_events.wait();
Ok(EventWaiter::new(rx, timeout.or(Some(30_000.0))))
}
#[tracing::instrument(level = "debug", skip_all, fields(guid = %self.guid()))]
pub async fn expect_close(&self, timeout: Option<f64>) -> Result<EventWaiter<()>> {
let rx = self.close_events.wait();
Ok(EventWaiter::new(rx, timeout.or(Some(30_000.0))))
}
#[tracing::instrument(level = "debug", skip_all, fields(guid = %self.guid()))]
pub async fn expect_console_message(
&self,
timeout: Option<f64>,
) -> Result<EventWaiter<crate::protocol::ConsoleMessage>> {
self.subscribe_if_idle(&self.console).await;
let rx = self.console.wait();
Ok(EventWaiter::new(rx, timeout.or(Some(30_000.0))))
}
#[tracing::instrument(level = "debug", skip_all, fields(guid = %self.guid()))]
pub async fn expect_event(
&self,
event: &str,
timeout: Option<f64>,
) -> crate::error::Result<EventWaiter<crate::protocol::EventValue>> {
let timeout_ms = timeout.or(Some(30_000.0));
let waiter = |rx| EventWaiter::new(rx, timeout_ms);
match event {
"page" => Ok(waiter(bridge(self.page_events.wait(), EventValue::Page))),
"close" => Ok(waiter(bridge(self.close_events.wait(), |()| {
EventValue::Close
}))),
"console" => {
self.subscribe_if_idle(&self.console).await;
Ok(waiter(bridge(
self.console.wait(),
EventValue::ConsoleMessage,
)))
}
"request" => {
self.subscribe_if_idle(&self.request).await;
Ok(waiter(bridge(self.request.wait(), EventValue::Request)))
}
"response" => {
self.subscribe_if_idle(&self.response).await;
Ok(waiter(bridge(self.response.wait(), EventValue::Response)))
}
"weberror" => Ok(waiter(bridge(self.weberror.wait(), EventValue::WebError))),
"serviceworker" => {
let (inner_tx, inner_rx) = oneshot::channel();
self.serviceworker_waiters.lock().unwrap().push(inner_tx);
Ok(waiter(bridge(inner_rx, EventValue::Worker)))
}
other => Err(crate::error::Error::InvalidArgument(format!(
"Unknown event name '{}'. Supported: page, close, console, request, response, \
weberror, serviceworker",
other
))),
}
}
}
type ForwardedHandler<T> = Arc<dyn Fn(T) -> CtxHandlerFuture + Send + Sync>;
type Forwarded<T> = Arc<Mutex<Vec<ForwardedHandler<T>>>>;
#[derive(Clone, Default)]
pub(super) struct Forwarders {
pub(super) download: Forwarded<Download>,
pub(super) frame_attached: Forwarded<Frame>,
pub(super) frame_detached: Forwarded<Frame>,
pub(super) frame_navigated: Forwarded<Frame>,
pub(super) page_load: Forwarded<Page>,
pub(super) page_close: Forwarded<Page>,
}
impl Forwarders {
pub(super) async fn wire(&self, page: &Page) {
_ = page.on_download(forward_to(&self.download)).await;
_ = page
.on_frameattached(forward_to(&self.frame_attached))
.await;
_ = page
.on_framedetached(forward_to(&self.frame_detached))
.await;
_ = page
.on_framenavigated(forward_to(&self.frame_navigated))
.await;
let (load, p) = (forward_to(&self.page_load), page.clone());
_ = page.on_load(move || load(p.clone())).await;
let (close, p) = (forward_to(&self.page_close), page.clone());
_ = page.on_close(move || close(p.clone())).await;
}
}
fn add_forwarded<T, F, Fut>(handlers: &Forwarded<T>, handler: F)
where
F: Fn(T) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
handlers
.lock()
.unwrap()
.push(Arc::new(move |value| Box::pin(handler(value))));
}
fn forward_to<T: Clone + Send + 'static>(
handlers: &Forwarded<T>,
) -> impl Fn(T) -> CtxHandlerFuture + Send + Sync + 'static {
let handlers = handlers.clone();
move |value| {
let handlers = handlers.clone();
Box::pin(async move {
run_all(&handlers, value).await;
Ok(())
})
}
}
async fn run_all<T: Clone>(handlers: &Forwarded<T>, value: T) {
let handlers = handlers.lock().unwrap().clone();
for handler in handlers {
if let Err(e) = handler(value.clone()).await {
tracing::warn!("context handler error: {e}");
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};
#[tokio::test]
async fn run_all_runs_every_handler_in_order_past_a_failure() {
let handlers: Forwarded<u8> = Forwarded::default();
let seen = Arc::new(Mutex::new(Vec::new()));
let calls = Arc::new(AtomicUsize::new(0));
for i in 0..3u8 {
let (seen, calls) = (seen.clone(), calls.clone());
add_forwarded(&handlers, move |v: u8| {
let (seen, calls) = (seen.clone(), calls.clone());
async move {
seen.lock().unwrap().push((i, v));
calls.fetch_add(1, Ordering::SeqCst);
if i == 1 {
Err(crate::error::Error::ProtocolError("boom".into()))
} else {
Ok(())
}
}
});
}
run_all(&handlers, 9).await;
assert_eq!(*seen.lock().unwrap(), vec![(0, 9), (1, 9), (2, 9)]);
assert_eq!(calls.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn forward_to_reaches_handlers_added_after_wiring() {
let handlers: Forwarded<u8> = Forwarded::default();
let page_side = forward_to(&handlers);
let seen = Arc::new(Mutex::new(Vec::new()));
let s = seen.clone();
add_forwarded(&handlers, move |v: u8| {
let s = s.clone();
async move {
s.lock().unwrap().push(v);
Ok(())
}
});
page_side(4).await.unwrap();
assert_eq!(*seen.lock().unwrap(), vec![4]);
}
}