use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{RwLock, oneshot};
use tracing::debug;
use viewpoint_cdp::CdpConnection;
use crate::error::PageError;
use crate::page::Page;
pub type PopupEventHandler =
Box<dyn Fn(Page) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>;
pub struct PopupManager {
handler: RwLock<Option<PopupEventHandler>>,
target_id: String,
connection: Arc<CdpConnection>,
session_id: String,
}
impl PopupManager {
pub fn new(connection: Arc<CdpConnection>, session_id: String, target_id: String) -> Self {
Self {
handler: RwLock::new(None),
target_id,
connection,
session_id,
}
}
pub async fn set_handler<F, Fut>(&self, handler: F)
where
F: Fn(Page) -> Fut + Send + Sync + 'static,
Fut: Future<Output = ()> + Send + 'static,
{
let boxed_handler: PopupEventHandler = Box::new(move |page| Box::pin(handler(page)));
let mut h = self.handler.write().await;
*h = Some(boxed_handler);
}
pub async fn remove_handler(&self) {
let mut h = self.handler.write().await;
*h = None;
}
pub async fn emit(&self, popup: Page) {
let handler = self.handler.read().await;
if let Some(ref h) = *handler {
h(popup).await;
}
}
pub fn is_opener(&self, opener_id: &str) -> bool {
self.target_id == opener_id
}
}
pub struct WaitForPopupBuilder<'a, F, Fut>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<(), crate::error::LocatorError>>,
{
page: &'a Page,
action: Option<F>,
timeout: Duration,
}
impl Page {
pub async fn on_popup<F, Fut>(&self, handler: F)
where
F: Fn(Page) -> Fut + Send + Sync + 'static,
Fut: Future<Output = ()> + Send + 'static,
{
self.popup_manager.set_handler(handler).await;
}
pub async fn off_popup(&self) {
self.popup_manager.remove_handler().await;
}
pub fn wait_for_popup<F, Fut>(&self, action: F) -> WaitForPopupBuilder<'_, F, Fut>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<(), crate::error::LocatorError>>,
{
WaitForPopupBuilder::new(self, action)
}
pub fn opener(&self) -> Option<&str> {
self.opener_target_id.as_deref()
}
}
impl<'a, F, Fut> WaitForPopupBuilder<'a, F, Fut>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<(), crate::error::LocatorError>>,
{
pub fn new(page: &'a Page, action: F) -> Self {
Self {
page,
action: Some(action),
timeout: Duration::from_secs(30),
}
}
#[must_use]
pub fn timeout(mut self, timeout: Duration) -> Self {
self.timeout = timeout;
self
}
pub async fn wait(mut self) -> Result<Page, PageError> {
use viewpoint_cdp::protocol::target_domain::{
AttachToTargetParams, AttachToTargetResult, TargetCreatedEvent,
};
let connection = self.page.connection().clone();
let target_id = self.page.target_id().to_string();
let _session_id = self.page.session_id().to_string();
let (tx, rx) = oneshot::channel::<Page>();
let tx = Arc::new(tokio::sync::Mutex::new(Some(tx)));
let mut events = connection.subscribe_events();
let tx_clone = tx.clone();
let connection_clone = connection.clone();
let target_id_clone = target_id.clone();
let popup_listener = tokio::spawn(async move {
while let Ok(event) = events.recv().await {
if event.method == "Target.targetCreated" {
if let Some(params) = &event.params {
if let Ok(created_event) =
serde_json::from_value::<TargetCreatedEvent>(params.clone())
{
let info = &created_event.target_info;
if info.target_type == "page"
&& info.opener_id.as_deref() == Some(&target_id_clone)
{
debug!("Popup detected: {}", info.target_id);
let attach_result: Result<AttachToTargetResult, _> =
connection_clone
.send_command(
"Target.attachToTarget",
Some(AttachToTargetParams {
target_id: info.target_id.clone(),
flatten: Some(true),
}),
None,
)
.await;
if let Ok(attach) = attach_result {
let popup_session = &attach.session_id;
let _ = connection_clone
.send_command::<(), serde_json::Value>(
"Page.enable",
None,
Some(popup_session),
)
.await;
let _ = connection_clone
.send_command::<(), serde_json::Value>(
"Network.enable",
None,
Some(popup_session),
)
.await;
let _ = connection_clone
.send_command::<(), serde_json::Value>(
"Runtime.enable",
None,
Some(popup_session),
)
.await;
let frame_tree: Result<
viewpoint_cdp::protocol::page::GetFrameTreeResult,
_,
> = connection_clone
.send_command(
"Page.getFrameTree",
None::<()>,
Some(popup_session),
)
.await;
if let Ok(tree) = frame_tree {
let frame_id = tree.frame_tree.frame.id;
let popup = Page::new(
connection_clone.clone(),
info.target_id.clone(),
attach.session_id.clone(),
frame_id,
);
let mut guard = tx_clone.lock().await;
if let Some(sender) = guard.take() {
let _ = sender.send(popup);
return;
}
}
}
}
}
}
}
}
});
let action = self.action.take().expect("action already consumed");
let action_result = action().await;
let result = match action_result {
Ok(()) => match tokio::time::timeout(self.timeout, rx).await {
Ok(Ok(popup)) => Ok(popup),
Ok(Err(_)) => Err(PageError::EvaluationFailed(
"Popup channel closed unexpectedly".to_string(),
)),
Err(_) => Err(PageError::EvaluationFailed(format!(
"wait_for_popup timed out after {:?}",
self.timeout
))),
},
Err(e) => Err(PageError::EvaluationFailed(e.to_string())),
};
popup_listener.abort();
result
}
}