use std::path::{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;
use crate::session::{liveness, Liveness};
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;
const INCOMING_PREFIX: &str = ".incoming-";
const COLLECT_CAP: usize = 64;
pub async fn arm(client: &CdpClient) -> Result<Armed, crate::BoxError> {
let tmp = tmp_root()?;
let _ = collect_abandoned(&tmp, COLLECT_CAP);
let dir = incoming_dir(&tmp);
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;
}
}
}
pub fn collect_abandoned(tmp: &Path, cap: usize) -> Vec<String> {
let Ok(entries) = std::fs::read_dir(tmp) else {
return Vec::new();
};
let mut names: Vec<String> = entries
.filter_map(|entry| entry.ok()?.file_name().into_string().ok())
.filter(|name| name.starts_with(INCOMING_PREFIX))
.collect();
names.sort_unstable();
names.truncate(cap);
let mut removed = Vec::new();
for name in names {
if !is_abandoned(&name) {
continue;
}
if std::fs::remove_dir_all(tmp.join(&name)).is_ok() {
removed.push(name);
}
}
removed
}
fn is_abandoned(name: &str) -> bool {
let Some(rest) = name.strip_prefix(INCOMING_PREFIX) else {
return false;
};
let Some((pid, _nanos)) = rest.split_once('-') else {
return false;
};
pid.parse::<u32>().is_ok_and(|pid| liveness(pid) == Liveness::Dead)
}
fn tmp_root() -> Result<PathBuf, crate::BoxError> {
let home = dirs::home_dir().ok_or("Could not determine home directory")?;
Ok(home.join(".chrome-agent").join("tmp"))
}
fn incoming_dir(tmp: &Path) -> PathBuf {
let nanos = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or_default();
tmp.join(format!("{INCOMING_PREFIX}{}-{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 tmp = scratch("incoming-names");
let first = incoming_dir(&tmp);
std::thread::sleep(Duration::from_millis(2));
let second = incoming_dir(&tmp);
assert_ne!(first, second);
assert!(first.to_string_lossy().contains(INCOMING_PREFIX));
std::fs::remove_dir_all(&tmp).ok();
}
fn scratch(tag: &str) -> PathBuf {
let nanos = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or_default();
let dir = std::env::temp_dir()
.join(format!("chrome-agent-{tag}-{}-{nanos}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
dir
}
fn transfer_dir(tmp: &Path, name: &str) -> PathBuf {
let dir = tmp.join(name);
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(dir.join("6f1c1f0e-guid"), b"partial").unwrap();
dir
}
#[cfg(unix)]
fn a_reaped_pid() -> u32 {
let mut child =
std::process::Command::new("/bin/sh").args(["-c", "exit 0"]).spawn().expect("spawn");
let pid = child.id();
child.wait().expect("wait");
pid
}
#[cfg(unix)]
#[test]
fn a_transfer_directory_is_collected_once_its_process_is_gone() {
let tmp = scratch("collect-abandoned");
let dead = a_reaped_pid();
assert_eq!(
liveness(dead),
Liveness::Dead,
"the pid was recycled between the wait and the probe, so this proves nothing"
);
let abandoned = transfer_dir(&tmp, &format!("{INCOMING_PREFIX}{dead}-1788086042802162000"));
let live = transfer_dir(
&tmp,
&format!("{INCOMING_PREFIX}{}-1788086042802162001", std::process::id()),
);
let removed = collect_abandoned(&tmp, COLLECT_CAP);
assert!(!abandoned.exists(), "a directory nothing can write to any more was kept");
assert_eq!(removed.len(), 1, "{removed:?}");
assert!(
live.exists(),
"a running process's transfer directory was taken, which on a concurrent agent is \
its download"
);
std::fs::remove_dir_all(&tmp).ok();
}
#[test]
fn an_unreadable_owner_or_another_command_s_file_is_left_alone() {
let tmp = scratch("collect-keeps");
let unparseable = transfer_dir(&tmp, &format!("{INCOMING_PREFIX}not-a-pid"));
let no_separator = transfer_dir(&tmp, INCOMING_PREFIX.trim_end_matches('-'));
let neighbour = tmp.join("shot-1788086042.png");
std::fs::write(&neighbour, b"png").unwrap();
assert!(collect_abandoned(&tmp, COLLECT_CAP).is_empty());
assert!(unparseable.exists());
assert!(no_separator.exists());
assert!(neighbour.exists());
std::fs::remove_dir_all(&tmp).ok();
}
#[cfg(unix)]
#[test]
fn the_cap_bounds_one_arming_and_the_backlog_still_converges() {
let tmp = scratch("collect-cap");
let dead = a_reaped_pid();
assert_eq!(liveness(dead), Liveness::Dead, "the pid was recycled");
for n in 0..5 {
transfer_dir(&tmp, &format!("{INCOMING_PREFIX}{dead}-178808604280216200{n}"));
}
assert_eq!(collect_abandoned(&tmp, 2).len(), 2, "the cap is not applied");
assert_eq!(collect_abandoned(&tmp, 2).len(), 2);
assert_eq!(collect_abandoned(&tmp, 2).len(), 1);
assert!(collect_abandoned(&tmp, 2).is_empty(), "the backlog did not drain");
std::fs::remove_dir_all(&tmp).ok();
}
}