use std::path::PathBuf;
use std::time::{Duration, Instant, SystemTime};
use serde_json::{json, Value};
use tokio::sync::broadcast;
use crate::cdp::client::CdpClient;
use crate::cdp::types::CdpEvent;
pub enum Transfer {
NeverBegan { waited_ms: u64 },
Completed { began: Began, bytes: u64, temp_path: PathBuf },
Canceled { began: Began, why: Cancelled },
Unfinished { began: Began, received: u64, total: u64, waited_ms: u64 },
}
#[derive(Clone, Copy, PartialEq, Eq)]
pub enum Cancelled {
ExceededCap,
ByBrowser,
}
pub struct Began {
pub guid: String,
pub suggested_filename: String,
pub url: String,
}
pub struct Armed {
events: broadcast::Receiver<CdpEvent>,
dir: PathBuf,
}
const SWEEP_ATTEMPTS: u32 = 5;
const SWEEP_GAP_MS: u64 = 30;
pub async fn arm(client: &CdpClient) -> Result<Armed, crate::BoxError> {
let dir = incoming_dir()?;
std::fs::create_dir_all(&dir)?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let _ = std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o700));
}
let events = client.events();
let path = dir.display().to_string();
if let Err(error) = client
.call::<_, Value>(
"Browser.setDownloadBehavior",
json!({"behavior": "allowAndName", "downloadPath": path, "eventsEnabled": true}),
)
.await
{
let _ = std::fs::remove_dir_all(&dir);
return Err(format!(
"download: Chrome refused to hand downloads to this session ({error}), so the click \
was not dispatched"
)
.into());
}
Ok(Armed { events, dir })
}
pub async fn disarm(client: &CdpClient) {
let _ = client
.call::<_, Value>("Browser.setDownloadBehavior", json!({"behavior": "default"}))
.await;
}
pub async fn collect(
client: &CdpClient,
armed: &mut Armed,
timeout: Duration,
max_bytes: u64,
) -> Transfer {
let started = Instant::now();
let mut began: Option<Began> = None;
let mut last_received = 0_u64;
let mut last_total = 0_u64;
let mut cancelled_by_us = false;
while let Some(left) = timeout.checked_sub(started.elapsed()) {
let event = match tokio::time::timeout(left, armed.events.recv()).await {
Ok(Err(broadcast::error::RecvError::Lagged(_))) => continue,
Err(_) | Ok(Err(broadcast::error::RecvError::Closed)) => break,
Ok(Ok(event)) => event,
};
match event.method.as_str() {
"Browser.downloadWillBegin" if began.is_none() => {
began = Some(Began {
guid: string_field(&event.params, "guid"),
suggested_filename: string_field(&event.params, "suggestedFilename"),
url: string_field(&event.params, "url"),
});
}
"Browser.downloadProgress" => {
let Some(current) = began.as_ref() else { continue };
if string_field(&event.params, "guid") != current.guid {
continue;
}
last_received = number_field(&event.params, "receivedBytes");
last_total = number_field(&event.params, "totalBytes");
let state = string_field(&event.params, "state");
if !cancelled_by_us && last_received.max(last_total) > max_bytes {
cancelled_by_us = true;
let _ = client
.call::<_, Value>(
"Browser.cancelDownload",
json!({"guid": current.guid}),
)
.await;
continue;
}
match state.as_str() {
"completed" => {
let began = began.take().expect("guarded above");
let temp_path = event
.params
.get("filePath")
.and_then(Value::as_str)
.map_or_else(|| armed.dir.join(&began.guid), PathBuf::from);
if cancelled_by_us {
let _ = std::fs::remove_file(&temp_path);
return Transfer::Canceled { began, why: Cancelled::ExceededCap };
}
return Transfer::Completed { began, bytes: last_received, temp_path };
}
"canceled" => {
let began = began.take().expect("guarded above");
let why = if cancelled_by_us {
Cancelled::ExceededCap
} else {
Cancelled::ByBrowser
};
return Transfer::Canceled { began, why };
}
_ => {}
}
}
_ => {}
}
}
let waited_ms = elapsed_ms(started);
match began {
None => Transfer::NeverBegan { waited_ms },
Some(began) => {
Transfer::Unfinished { began, received: last_received, total: last_total, waited_ms }
}
}
}
pub fn place(
completed_path: &std::path::Path,
suggested: &str,
out: Option<&str>,
) -> Result<(String, u64), crate::BoxError> {
let destination = super::download::resolve_named_path(out, suggested)?;
if let Some(parent) = destination.parent() {
std::fs::create_dir_all(parent)?;
}
if std::fs::rename(completed_path, &destination).is_err() {
std::fs::copy(completed_path, &destination)?;
let _ = std::fs::remove_file(completed_path);
}
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let _ = std::fs::set_permissions(&destination, std::fs::Permissions::from_mode(0o600));
}
let bytes = std::fs::metadata(&destination)?.len();
Ok((destination.display().to_string(), bytes))
}
pub async fn clean_up(armed: &Armed) {
for attempt in 0..SWEEP_ATTEMPTS {
let _ = std::fs::remove_dir_all(&armed.dir);
if !armed.dir.exists() {
return;
}
if attempt + 1 < SWEEP_ATTEMPTS {
tokio::time::sleep(Duration::from_millis(SWEEP_GAP_MS)).await;
}
}
}
fn incoming_dir() -> Result<PathBuf, crate::BoxError> {
let nanos = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or_default();
let home = dirs::home_dir().ok_or("Could not determine home directory")?;
Ok(home
.join(".chrome-agent")
.join("tmp")
.join(format!(".incoming-{}-{nanos}", std::process::id())))
}
fn string_field(params: &Value, key: &str) -> String {
params.get(key).and_then(Value::as_str).unwrap_or_default().to_string()
}
fn number_field(params: &Value, key: &str) -> u64 {
let Some(value) = params.get(key) else { return 0 };
if let Some(exact) = value.as_u64() {
return exact;
}
let truncated = value.as_f64().map_or(0_i64, |number| number.trunc() as i64);
u64::try_from(truncated).unwrap_or(0)
}
fn elapsed_ms(since: Instant) -> u64 {
u64::try_from(since.elapsed().as_millis()).unwrap_or(u64::MAX)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn byte_counters_survive_an_integer_or_a_float() {
assert_eq!(number_field(&json!({"totalBytes": 22}), "totalBytes"), 22);
assert_eq!(number_field(&json!({"totalBytes": 22.0}), "totalBytes"), 22);
assert_eq!(number_field(&json!({"totalBytes": -1}), "totalBytes"), 0);
assert_eq!(number_field(&json!({}), "totalBytes"), 0);
}
#[test]
fn a_missing_string_field_is_empty_not_a_panic() {
assert_eq!(string_field(&json!({}), "guid"), "");
assert_eq!(string_field(&json!({"guid": 7}), "guid"), "");
assert_eq!(string_field(&json!({"guid": "abc"}), "guid"), "abc");
}
#[test]
fn each_invocation_gets_its_own_incoming_directory() {
let first = incoming_dir().unwrap();
std::thread::sleep(Duration::from_millis(2));
let second = incoming_dir().unwrap();
assert_ne!(first, second);
assert!(first.to_string_lossy().contains(".incoming-"));
}
}