diff options
Diffstat (limited to 'src/diagnostics.rs')
| -rw-r--r-- | src/diagnostics.rs | 172 |
1 files changed, 172 insertions, 0 deletions
diff --git a/src/diagnostics.rs b/src/diagnostics.rs new file mode 100644 index 0000000..a2052ec --- /dev/null +++ b/src/diagnostics.rs @@ -0,0 +1,172 @@ +use crate::model::Settings; +use anyhow::{Result, anyhow, ensure}; +use futures_util::TryStreamExt; +use serde::{Deserialize, Serialize}; +use std::time::{Duration, Instant}; +use tokio::io::AsyncReadExt; + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct Observation { + pub configuration_id: String, + pub generation: u64, + pub kind: String, + pub success: bool, + pub elapsed_ms: u64, + pub bytes: u64, + pub bytes_per_second: Option<f64>, + pub stop_reason: String, + pub retry_after_seconds: Option<u64>, +} + +/// The caller supplies an authenticated, service-owned SOCKS endpoint. This client +/// never inherits environment proxies or creates a direct route to the target. +pub async fn measure( + settings: &Settings, + proxy: &str, + id: String, + generation: u64, + speed: bool, +) -> Result<Observation> { + settings.validate()?; + let started = Instant::now(); + let duration = if speed { + settings.maximum_duration_seconds + } else { + 5 + }; + let deadline = tokio::time::Instant::now() + Duration::from_secs(duration); + let cap = if speed { + settings.maximum_download_mib * 1_048_576 + } else { + 4096 + }; + let proxy = + reqwest::Proxy::all(proxy).map_err(|_| anyhow!("Invalid protected probe endpoint"))?; + let client = reqwest::Client::builder() + .no_proxy() + .proxy(proxy) + .redirect(reqwest::redirect::Policy::none()) + .timeout(Duration::from_secs(duration)) + .user_agent(concat!("Roentgen/", env!("CARGO_PKG_VERSION"))) + .build() + .map_err(|_| anyhow!("Could not initialize diagnostic TLS client"))?; + let url = if speed { + &settings.speed_url + } else { + &settings.ping_url + }; + let response = tokio::time::timeout_at(deadline, client.get(url).send()) + .await + .map_err(|_| anyhow!("Diagnostic request reached its deadline"))? + .map_err(|_| anyhow!("Proxied HTTPS request failed"))?; + if response.status().as_u16() == 429 || response.status().as_u16() == 503 { + let retry = response + .headers() + .get("Retry-After") + .and_then(|v| v.to_str().ok()) + .and_then(|v| { + v.parse::<u64>().ok().or_else(|| { + httpdate::parse_http_date(v) + .ok() + .and_then(|date| date.duration_since(std::time::SystemTime::now()).ok()) + .map(|d| d.as_secs().saturating_add(1)) + }) + }) + .unwrap_or(30) + .max(30); + return Ok(Observation { + configuration_id: id, + generation, + kind: if speed { "speed" } else { "ping" }.into(), + success: false, + elapsed_ms: started.elapsed().as_millis() as u64, + bytes: 0, + bytes_per_second: None, + stop_reason: "provider-rate-limit".into(), + retry_after_seconds: Some(retry), + }); + } + ensure!( + response.status().as_u16() == 200 || !speed && response.status().as_u16() == 204, + "Diagnostic endpoint returned an unexpected HTTP status" + ); + let content_type = response + .headers() + .get("Content-Type") + .and_then(|v| v.to_str().ok()) + .unwrap_or("") + .to_ascii_lowercase(); + ensure!( + !content_type.contains("html"), + "Diagnostic endpoint returned HTML" + ); + let status = response.status().as_u16(); + let stream = response.bytes_stream().map_err(std::io::Error::other); + let mut reader = tokio_util::io::StreamReader::new(stream); + let mut buffer = [0u8; 16384]; + let mut prefix = Vec::new(); + let mut bytes = 0; + let reason; + loop { + if speed && bytes == cap { + reason = "byte-limit"; + break; + } + let remaining = if speed { + (cap - bytes).min(buffer.len() as u64) as usize + } else { + buffer.len().min((cap + 1 - bytes) as usize) + }; + let chunk = + match tokio::time::timeout_at(deadline, reader.read(&mut buffer[..remaining])).await { + Err(_) if speed && bytes > 0 => { + reason = "duration-limit"; + break; + } + Err(_) => return Err(anyhow!("Diagnostic response timed out")), + Ok(Err(_)) => { + return Err(anyhow!( + "Diagnostic response failed before clean completion" + )); + } + Ok(Ok(value)) => value, + }; + if chunk == 0 { + reason = "complete"; + break; + } + if prefix.len() < 512 { + prefix.extend_from_slice(&buffer[..chunk.min(512 - prefix.len())]); + } + bytes += chunk as u64; + if !speed { + ensure!(bytes <= cap, "Ping body exceeds 4 KiB"); + } + } + let prefix = String::from_utf8_lossy(&prefix) + .trim_start() + .to_ascii_lowercase(); + ensure!( + !prefix.starts_with("<!doctype html") && !prefix.starts_with("<html"), + "Diagnostic endpoint returned an HTML document" + ); + ensure!(!speed || bytes > 0, "Speed response contained no data"); + if !speed && settings.ping_url == Settings::default().ping_url { + ensure!( + bytes == 0 && status == 200, + "Default Ping must complete with HTTP 200 and an empty body" + ); + } + let elapsed = started.elapsed(); + Ok(Observation { + configuration_id: id, + generation, + kind: if speed { "speed" } else { "ping" }.into(), + success: true, + elapsed_ms: elapsed.as_millis() as u64, + bytes, + bytes_per_second: speed.then(|| bytes as f64 / elapsed.as_secs_f64().max(0.000_001)), + stop_reason: reason.into(), + retry_after_seconds: None, + }) +} |