1use std::marker::PhantomData;
2use std::sync::Arc;
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::{
5 fs,
6 io::{Read, Write},
7};
8
9use crate::proto::context_session_command::Command as ContextCommand;
10use crate::proto::context_session_event::Event as ContextEvent;
11use crate::proto::hook_completed_event::Result as HookCompletionResult;
12use crate::proto::register_hook_command::Hook as RegisterHook;
13use crate::proto::{
14 ContextSessionCommand, ReadFileChunkCommand, RegisterDownloadHook, RegisterFileChooserHook,
15 RegisterHookCommand, RegisterNewPageHook, SaveDownloadCommand, SetFileChooserFilesCommand,
16 UploadFileChunkCommand, WaitForHookCommand,
17};
18use std::path::Path;
19use tokio::sync::Mutex as AsyncMutex;
20
21use super::command::command_retry_options;
22use super::tab::ensure_tab_open;
23use super::types::{CommandOptions, Error, Result, Tab, TabInner, TabState};
24
25static TRANSFER_COUNTER: AtomicU64 = AtomicU64::new(1);
26
27pub trait HookType: private::Sealed + Clone + Send + Sync + 'static {
28 type Output;
29}
30
31mod private {
32 use super::*;
33
34 pub trait Sealed {
35 fn name() -> &'static str;
36 fn decode(page: &Tab, result: HookCompletionResult) -> Result<<Self as HookType>::Output>
37 where
38 Self: HookType;
39 }
40}
41
42#[derive(Debug, Clone, Copy, Default)]
43pub struct DialogHook;
44pub const DIALOG: DialogHook = DialogHook;
45impl HookType for DialogHook {
46 type Output = Dialog;
47}
48impl private::Sealed for DialogHook {
49 fn name() -> &'static str {
50 "dialog"
51 }
52 fn decode(page: &Tab, result: HookCompletionResult) -> Result<Dialog> {
53 let HookCompletionResult::Dialog(d) = result else {
54 return Err(Error::new("invalid dialog hook result"));
55 };
56 Ok(Dialog {
57 page: page.clone(),
58 id: d.dialog_id,
59 kind: d.r#type,
60 message: d.message,
61 default_value: d.default_value,
62 })
63 }
64}
65
66#[derive(Clone)]
67pub struct Dialog {
68 page: Tab,
69 pub id: String,
70 pub kind: String,
71 pub message: String,
72 pub default_value: String,
73}
74impl Dialog {
75 pub fn page(&self) -> &Tab {
76 &self.page
77 }
78 pub async fn accept(&self, prompt_text: Option<&str>) -> Result<()> {
79 self.accept_with_options(prompt_text, CommandOptions::default())
80 .await
81 }
82 pub async fn dismiss(&self) -> Result<()> {
83 self.dismiss_with_options(CommandOptions::default()).await
84 }
85 pub async fn accept_with_options(
86 &self,
87 prompt_text: Option<&str>,
88 options: CommandOptions,
89 ) -> Result<()> {
90 self.handle(true, prompt_text, options).await
91 }
92 pub async fn dismiss_with_options(&self, options: CommandOptions) -> Result<()> {
93 self.handle(false, None, options).await
94 }
95 async fn handle(
96 &self,
97 accept: bool,
98 prompt_text: Option<&str>,
99 options: CommandOptions,
100 ) -> Result<()> {
101 let mut state = self.page.inner.state.lock().await;
102 let handle = self.page.ensure_handle(&mut state).await?;
103 ensure_tab_open(handle, &self.page.inner.session_id)?;
104 handle
105 .command_tx
106 .send(ContextSessionCommand {
107 surface_session_id: self.page.inner.surface_session_id.clone(),
108 context_session_id: self.page.inner.session_id.clone(),
109 command: Some(ContextCommand::HandleDialog(
110 crate::proto::HandleDialogCommand {
111 dialog_id: self.id.clone(),
112 accept,
113 prompt_text: prompt_text.map(str::to_owned),
114 retry_options: command_retry_options(options.timeout_ms),
115 },
116 )),
117 })
118 .await
119 .map_err(|_| Error::new("failed to send HandleDialogCommand"))?;
120 loop {
121 let event = handle
122 .events
123 .message()
124 .await?
125 .ok_or_else(|| Error::new("page closed while handling dialog"))?;
126 match event.event {
127 Some(ContextEvent::DialogHandled(d)) if d.dialog_id == self.id => return Ok(()),
128 Some(ContextEvent::Error(e)) => return Err(Error::new(e.message)),
129 Some(ContextEvent::Closed(_)) => {
130 return Err(Error::new("page closed while handling dialog"));
131 }
132 _ => {}
133 }
134 }
135 }
136}
137
138#[derive(Debug, Clone, Copy, Default)]
139pub struct NewPage;
140
141pub const NEW_PAGE: NewPage = NewPage;
142
143impl HookType for NewPage {
144 type Output = Tab;
145}
146
147impl private::Sealed for NewPage {
148 fn name() -> &'static str {
149 "new_page"
150 }
151
152 fn decode(page: &Tab, result: HookCompletionResult) -> Result<<Self as HookType>::Output> {
153 let HookCompletionResult::NewPage(new_page) = result else {
154 return Err(Error::new("new page hook completed with an invalid result"));
155 };
156 Ok(Tab {
157 inner: Arc::new(TabInner {
158 runtime: Arc::clone(&page.inner.runtime),
159 surface_session_id: page.inner.surface_session_id.clone(),
160 session_id: new_page.context_session_id,
161 state: AsyncMutex::new(TabState::default()),
162 }),
163 })
164 }
165}
166
167#[derive(Debug, Clone, Copy, Default)]
168pub struct FileChooserHook;
169
170pub const FILE_CHOOSER: FileChooserHook = FileChooserHook;
171
172impl HookType for FileChooserHook {
173 type Output = FileChooser;
174}
175
176impl private::Sealed for FileChooserHook {
177 fn name() -> &'static str {
178 "file_chooser"
179 }
180
181 fn decode(page: &Tab, result: HookCompletionResult) -> Result<<Self as HookType>::Output> {
182 let HookCompletionResult::FileChooser(file_chooser) = result else {
183 return Err(Error::new(
184 "file chooser hook completed with an invalid result",
185 ));
186 };
187 Ok(FileChooser {
188 page: page.clone(),
189 id: file_chooser.file_chooser_id,
190 is_multiple: file_chooser.is_multiple,
191 })
192 }
193}
194
195#[derive(Clone)]
196pub struct FileChooser {
197 page: Tab,
198 id: String,
199 is_multiple: bool,
200}
201
202impl FileChooser {
203 pub fn id(&self) -> &str {
204 &self.id
205 }
206
207 pub fn is_multiple(&self) -> bool {
208 self.is_multiple
209 }
210
211 pub fn page(&self) -> &Tab {
212 &self.page
213 }
214
215 pub async fn set_file(&self, file: impl AsRef<Path>) -> Result<()> {
216 self.set_files([file]).await
217 }
218
219 pub async fn set_files<I, P>(&self, files: I) -> Result<()>
220 where
221 I: IntoIterator<Item = P>,
222 P: AsRef<Path>,
223 {
224 let files = files
225 .into_iter()
226 .map(|file| file.as_ref().to_path_buf())
227 .collect::<Vec<_>>();
228 let mut state = self.page.inner.state.lock().await;
229 let handle = self.page.ensure_handle(&mut state).await?;
230 ensure_tab_open(handle, &self.page.inner.session_id)?;
231 let mut file_ids = Vec::with_capacity(files.len());
232 for path in files {
233 let name = path
234 .file_name()
235 .and_then(|name| name.to_str())
236 .ok_or_else(|| Error::new("upload path requires a valid file name"))?
237 .to_string();
238 let transfer_id = format!(
239 "rust-upload-{}-{}",
240 std::process::id(),
241 TRANSFER_COUNTER.fetch_add(1, Ordering::Relaxed)
242 );
243 let mut source = fs::File::open(&path)
244 .map_err(|error| Error::new(format!("open upload {}: {error}", path.display())))?;
245 let size = source
246 .metadata()
247 .map_err(|error| Error::new(format!("inspect upload {}: {error}", path.display())))?
248 .len();
249 let mut offset = 0_u64;
250 loop {
251 let mut data = vec![0; 256 * 1024];
252 let count = source.read(&mut data).map_err(|error| {
253 Error::new(format!("read upload {}: {error}", path.display()))
254 })?;
255 data.truncate(count);
256 let last = offset + count as u64 >= size;
257 handle
258 .command_tx
259 .send(ContextSessionCommand {
260 surface_session_id: self.page.inner.surface_session_id.clone(),
261 context_session_id: self.page.inner.session_id.clone(),
262 command: Some(ContextCommand::UploadFileChunk(UploadFileChunkCommand {
263 transfer_id: transfer_id.clone(),
264 name: name.clone(),
265 offset,
266 data,
267 last,
268 })),
269 })
270 .await
271 .map_err(|_| Error::new("failed to send UploadFileChunkCommand"))?;
272 offset += count as u64;
273 if last {
274 break;
275 }
276 }
277 loop {
278 let event = handle.events.message().await?.ok_or_else(|| {
279 Error::new("page session closed while uploading chooser file")
280 })?;
281 match event.event {
282 Some(ContextEvent::FileUploaded(result))
283 if result.transfer_id == transfer_id =>
284 {
285 file_ids.push(result.file_id);
286 break;
287 }
288 Some(ContextEvent::Error(error)) => {
289 return Err(Error::new(format!(
290 "page session error while uploading chooser file: {}",
291 error.message
292 )));
293 }
294 _ => {}
295 }
296 }
297 }
298 handle
299 .command_tx
300 .send(ContextSessionCommand {
301 surface_session_id: self.page.inner.surface_session_id.clone(),
302 context_session_id: self.page.inner.session_id.clone(),
303 command: Some(ContextCommand::SetFileChooserFiles(
304 SetFileChooserFilesCommand {
305 file_chooser_id: self.id.clone(),
306 file_ids,
307 retry_options: None,
308 },
309 )),
310 })
311 .await
312 .map_err(|_| Error::new("failed to send SetFileChooserFilesCommand"))?;
313 loop {
314 let event = handle
315 .events
316 .message()
317 .await?
318 .ok_or_else(|| Error::new("page session closed while setting chooser files"))?;
319 match event.event {
320 Some(ContextEvent::FileChooserFilesSet(result))
321 if result.file_chooser_id == self.id =>
322 {
323 return Ok(());
324 }
325 Some(ContextEvent::Error(error)) => {
326 return Err(Error::new(format!(
327 "page session error while setting chooser files: {}",
328 error.message
329 )));
330 }
331 _ => {}
332 }
333 }
334 }
335}
336
337#[derive(Debug, Clone, Copy, Default)]
338pub struct DownloadHook;
339
340pub const DOWNLOAD: DownloadHook = DownloadHook;
341
342impl HookType for DownloadHook {
343 type Output = Download;
344}
345
346impl private::Sealed for DownloadHook {
347 fn name() -> &'static str {
348 "download"
349 }
350
351 fn decode(page: &Tab, result: HookCompletionResult) -> Result<<Self as HookType>::Output> {
352 let HookCompletionResult::Download(download) = result else {
353 return Err(Error::new("download hook completed with an invalid result"));
354 };
355 Ok(Download {
356 page: page.clone(),
357 id: download.download_id,
358 url: download.url,
359 suggested_filename: download.suggested_filename,
360 })
361 }
362}
363
364#[derive(Clone)]
365pub struct Download {
366 page: Tab,
367 id: String,
368 url: String,
369 suggested_filename: String,
370}
371
372impl Download {
373 pub fn id(&self) -> &str {
374 &self.id
375 }
376
377 pub fn page(&self) -> &Tab {
378 &self.page
379 }
380
381 pub fn url(&self) -> &str {
382 &self.url
383 }
384
385 pub fn suggested_filename(&self) -> &str {
386 &self.suggested_filename
387 }
388
389 pub async fn save_as(&self, path: impl AsRef<Path>) -> Result<()> {
390 self.save_as_with_options(path, CommandOptions::default())
391 .await
392 }
393
394 pub async fn save_as_with_options(
395 &self,
396 path: impl AsRef<Path>,
397 options: CommandOptions,
398 ) -> Result<()> {
399 let mut state = self.page.inner.state.lock().await;
400 let handle = self.page.ensure_handle(&mut state).await?;
401 ensure_tab_open(handle, &self.page.inner.session_id)?;
402 handle
403 .command_tx
404 .send(ContextSessionCommand {
405 surface_session_id: self.page.inner.surface_session_id.clone(),
406 context_session_id: self.page.inner.session_id.clone(),
407 command: Some(ContextCommand::SaveDownload(SaveDownloadCommand {
408 download_id: self.id.clone(),
409 retry_options: command_retry_options(options.timeout_ms),
410 })),
411 })
412 .await
413 .map_err(|_| Error::new("failed to send SaveDownloadCommand"))?;
414 loop {
415 let event = handle
416 .events
417 .message()
418 .await?
419 .ok_or_else(|| Error::new("page session closed while saving download"))?;
420 match event.event {
421 Some(ContextEvent::DownloadSaved(result)) if result.download_id == self.id => {
422 let destination = path.as_ref();
423 let temporary = destination.with_extension(format!(
424 "allwright-{}.tmp",
425 TRANSFER_COUNTER.fetch_add(1, Ordering::Relaxed)
426 ));
427 let mut output = fs::OpenOptions::new()
428 .write(true)
429 .create_new(true)
430 .open(&temporary)
431 .map_err(|error| Error::new(format!("create download file: {error}")))?;
432 let mut offset = 0_u64;
433 loop {
434 handle
435 .command_tx
436 .send(ContextSessionCommand {
437 surface_session_id: self.page.inner.surface_session_id.clone(),
438 context_session_id: self.page.inner.session_id.clone(),
439 command: Some(ContextCommand::ReadFileChunk(
440 ReadFileChunkCommand {
441 file_id: result.file_id.clone(),
442 offset,
443 max_bytes: 256 * 1024,
444 },
445 )),
446 })
447 .await
448 .map_err(|_| Error::new("failed to send ReadFileChunkCommand"))?;
449 let chunk = handle.events.message().await?.ok_or_else(|| {
450 Error::new("page session closed while downloading file")
451 })?;
452 match chunk.event {
453 Some(ContextEvent::FileChunk(chunk))
454 if chunk.file_id == result.file_id && chunk.offset == offset =>
455 {
456 output.write_all(&chunk.data).map_err(|error| {
457 Error::new(format!("write download file: {error}"))
458 })?;
459 offset += chunk.data.len() as u64;
460 if chunk.last {
461 drop(output);
462 fs::rename(&temporary, destination).map_err(|error| {
463 Error::new(format!("finish download file: {error}"))
464 })?;
465 return Ok(());
466 }
467 }
468 Some(ContextEvent::Error(error)) => {
469 let _ = fs::remove_file(&temporary);
470 return Err(Error::new(format!(
471 "page session error while downloading file: {}",
472 error.message
473 )));
474 }
475 _ => {}
476 }
477 }
478 }
479 Some(ContextEvent::Error(error)) => {
480 return Err(Error::new(format!(
481 "page session error while saving download: {}",
482 error.message
483 )));
484 }
485 _ => {}
486 }
487 }
488 }
489}
490
491pub struct Hook<T: HookType> {
492 page: Tab,
493 id: String,
494 _type: PhantomData<T>,
495}
496
497impl<T: HookType> Hook<T> {
498 pub fn id(&self) -> &str {
499 &self.id
500 }
501
502 pub async fn wait(&self) -> Result<T::Output> {
503 self.wait_with_options(CommandOptions::default()).await
504 }
505
506 pub async fn wait_with_options(&self, options: CommandOptions) -> Result<T::Output> {
507 let mut state = self.page.inner.state.lock().await;
508 let handle = self.page.ensure_handle(&mut state).await?;
509 ensure_tab_open(handle, &self.page.inner.session_id)?;
510 handle
511 .command_tx
512 .send(ContextSessionCommand {
513 surface_session_id: self.page.inner.surface_session_id.clone(),
514 context_session_id: self.page.inner.session_id.clone(),
515 command: Some(ContextCommand::WaitForHook(WaitForHookCommand {
516 hook_id: self.id.clone(),
517 retry_options: command_retry_options(options.timeout_ms),
518 })),
519 })
520 .await
521 .map_err(|_| Error::new("failed to send WaitForHookCommand"))?;
522
523 loop {
524 let event = handle
525 .events
526 .message()
527 .await?
528 .ok_or_else(|| Error::new("page session closed while waiting for hook"))?;
529 match event.event {
530 Some(ContextEvent::HookCompleted(completed)) if completed.hook_id == self.id => {
531 let result = completed
532 .result
533 .ok_or_else(|| Error::new("hook completed without a result"))?;
534 return T::decode(&self.page, result);
535 }
536 Some(ContextEvent::Error(error)) => {
537 return Err(Error::new(format!(
538 "page session error while waiting for hook: {}",
539 error.message
540 )));
541 }
542 _ => {}
543 }
544 }
545 }
546}
547
548impl Tab {
549 pub async fn register_hook<T: HookType>(&self, _hook_type: T) -> Result<Hook<T>> {
550 let mut state = self.inner.state.lock().await;
551 let handle = self.ensure_handle(&mut state).await?;
552 ensure_tab_open(handle, &self.inner.session_id)?;
553 handle
554 .command_tx
555 .send(ContextSessionCommand {
556 surface_session_id: self.inner.surface_session_id.clone(),
557 context_session_id: self.inner.session_id.clone(),
558 command: Some(ContextCommand::RegisterHook(RegisterHookCommand {
559 hook: match T::name() {
560 "new_page" => Some(RegisterHook::NewPage(RegisterNewPageHook {})),
561 "dialog" => Some(RegisterHook::Dialog(crate::proto::RegisterDialogHook {})),
562 "file_chooser" => {
563 Some(RegisterHook::FileChooser(RegisterFileChooserHook {}))
564 }
565 "download" => Some(RegisterHook::Download(RegisterDownloadHook {})),
566 _ => return Err(Error::new("unsupported hook type")),
567 },
568 })),
569 })
570 .await
571 .map_err(|_| Error::new("failed to send RegisterHookCommand"))?;
572
573 loop {
574 let event = handle
575 .events
576 .message()
577 .await?
578 .ok_or_else(|| Error::new("page session closed while registering hook"))?;
579 match event.event {
580 Some(ContextEvent::HookRegistered(registered)) => {
581 return Ok(Hook {
582 page: self.clone(),
583 id: registered.hook_id,
584 _type: PhantomData,
585 });
586 }
587 Some(ContextEvent::Error(error)) => {
588 return Err(Error::new(format!(
589 "page session error while registering hook: {}",
590 error.message
591 )));
592 }
593 _ => {}
594 }
595 }
596 }
597}