millipede_browser/
hooks.rs1use std::{fmt, sync::Arc};
4
5use futures_util::future::BoxFuture;
6
7use crate::{BrowserError, BrowserPage, LaunchContext, PageId, PageOptions};
8
9pub type PreLaunchHook = Arc<dyn Fn(&mut LaunchContext) + Send + Sync>;
11
12pub type PagePrepHook = Arc<dyn Fn(&mut PageOptions) + Send + Sync>;
14
15pub type PageHook = Arc<
17 dyn for<'a> Fn(&'a dyn BrowserPage, &'a PageOptions) -> BoxFuture<'a, Result<(), BrowserError>>
18 + Send
19 + Sync,
20>;
21
22pub type PageClosedHook = Arc<dyn Fn(PageId) + Send + Sync>;
24
25#[derive(Clone, Default)]
33#[must_use = "browser hooks do nothing unless installed on a browser builder"]
34pub struct BrowserHooks {
35 pub pre_launch: Vec<PreLaunchHook>,
37 pub pre_page_create: Vec<PagePrepHook>,
39 pub post_page_create: Vec<PageHook>,
41 pub pre_page_close: Vec<PageHook>,
43 pub post_page_close: Vec<PageClosedHook>,
45}
46
47impl BrowserHooks {
48 pub fn defaults() -> Self {
50 Self::default().with_session_cookie_sync()
51 }
52
53 pub fn push_pre_launch(
55 mut self,
56 hook: impl Fn(&mut LaunchContext) + Send + Sync + 'static,
57 ) -> Self {
58 self.pre_launch.push(Arc::new(hook));
59 self
60 }
61
62 pub fn push_pre_page_create(
64 mut self,
65 hook: impl Fn(&mut PageOptions) + Send + Sync + 'static,
66 ) -> Self {
67 self.pre_page_create.push(Arc::new(hook));
68 self
69 }
70
71 pub fn push_post_page_create<F>(mut self, hook: F) -> Self
73 where
74 F: for<'a> Fn(
75 &'a dyn BrowserPage,
76 &'a PageOptions,
77 ) -> BoxFuture<'a, Result<(), BrowserError>>
78 + Send
79 + Sync
80 + 'static,
81 {
82 self.post_page_create.push(Arc::new(hook));
83 self
84 }
85
86 pub fn push_pre_page_close<F>(mut self, hook: F) -> Self
88 where
89 F: for<'a> Fn(
90 &'a dyn BrowserPage,
91 &'a PageOptions,
92 ) -> BoxFuture<'a, Result<(), BrowserError>>
93 + Send
94 + Sync
95 + 'static,
96 {
97 self.pre_page_close.push(Arc::new(hook));
98 self
99 }
100
101 pub fn push_post_page_close(mut self, hook: impl Fn(PageId) + Send + Sync + 'static) -> Self {
103 self.post_page_close.push(Arc::new(hook));
104 self
105 }
106
107 pub fn with_launch_args(self, args: Vec<String>) -> Self {
109 self.push_pre_launch(move |ctx| ctx.extra_args.extend(args.iter().cloned()))
110 }
111
112 pub fn with_session_cookie_sync(self) -> Self {
116 self.push_post_page_create(|page, opts| {
117 Box::pin(async move {
118 if let Some(session) = &opts.session {
119 let cookies = session.cookie_jar().export_cookies();
120 if !cookies.is_empty() {
121 page.set_cookies(&cookies).await?;
122 }
123 }
124 if !opts.extra_headers.is_empty() {
125 page.set_extra_headers(&opts.extra_headers).await?;
126 }
127 Ok(())
128 })
129 })
130 .push_pre_page_close(|page, opts| {
131 Box::pin(async move {
132 if let Some(session) = &opts.session {
133 let cookies = page.cookies().await?;
134 session.cookie_jar().import_cookies(&cookies);
135 tracing::debug!(cookie_count = cookies.len(), "synchronized browser cookies");
136 }
137 Ok(())
138 })
139 })
140 }
141
142 pub fn with_fingerprint(
147 self,
148 generator: Arc<millipede_fingerprint::BrowserFingerprintGenerator>,
149 ) -> Self {
150 self.push_post_page_create(move |page, opts| {
151 let generator = Arc::clone(&generator);
152 Box::pin(async move {
153 let seed = opts
154 .session
155 .as_ref()
156 .map(|session| session.id().as_str().to_owned())
157 .unwrap_or_else(|| "anonymous".to_owned());
158 let profile = generator.generate(&seed);
159 let mut headers = http::HeaderMap::new();
160 if let Ok(user_agent) = http::HeaderValue::from_str(&profile.user_agent) {
161 headers.insert(http::header::USER_AGENT, user_agent);
162 }
163 for (name, value) in profile.headers {
164 if let Ok(name) = http::HeaderName::from_bytes(name.as_bytes()) {
165 if let Ok(value) = http::HeaderValue::from_str(&value) {
166 headers.insert(name, value);
167 }
168 }
169 }
170 page.set_extra_headers(&headers).await?;
171 Ok(())
172 })
173 })
174 }
175}
176
177impl fmt::Debug for BrowserHooks {
178 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
179 formatter
180 .debug_struct("BrowserHooks")
181 .field("pre_launch_count", &self.pre_launch.len())
182 .field("pre_page_create_count", &self.pre_page_create.len())
183 .field("post_page_create_count", &self.post_page_create.len())
184 .field("pre_page_close_count", &self.pre_page_close.len())
185 .field("post_page_close_count", &self.post_page_close.len())
186 .finish()
187 }
188}