diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/diagnostics.rs | 31 | ||||
| -rw-r--r-- | src/diagnostics/error_document.rs | 170 |
2 files changed, 190 insertions, 11 deletions
diff --git a/src/diagnostics.rs b/src/diagnostics.rs index f6ed82a..d8f8620 100644 --- a/src/diagnostics.rs +++ b/src/diagnostics.rs @@ -1,3 +1,5 @@ +mod error_document; + use crate::model::Settings; use anyhow::{Result, anyhow, ensure}; use futures_util::TryStreamExt; @@ -113,6 +115,7 @@ pub async fn measure( let mut reader = tokio_util::io::StreamReader::new(stream); let mut buffer = [0u8; 16384]; let mut prefix = Vec::new(); + let mut error_document = error_document::ErrorDocument::default(); let mut bytes = 0; let reason; loop { @@ -144,8 +147,23 @@ pub async fn measure( break; } if prefix.len() < 512 { - prefix.extend_from_slice(&buffer[..chunk.min(512 - prefix.len())]); + let mut inspect = &buffer[..chunk]; + if prefix.is_empty() { + inspect = inspect.trim_ascii_start(); + while inspect + .first() + .is_some_and(|byte| matches!(byte, 0xef | 0xbb | 0xbf)) + { + inspect = inspect[1..].trim_ascii_start(); + } + } + prefix.extend_from_slice(&inspect[..inspect.len().min(512 - prefix.len())]); } + error_document.feed(&buffer[..chunk]); + ensure!( + !error_document.confirmed_error(), + "Diagnostic endpoint returned an error document" + ); bytes += chunk as u64; if !speed { ensure!(bytes <= cap, "Ping body exceeds 4 KiB"); @@ -158,17 +176,8 @@ pub async fn measure( !prefix.starts_with("<!doctype html") && !prefix.starts_with("<html"), "Diagnostic endpoint returned an HTML document" ); - let json_error = serde_json::from_str::<serde_json::Value>(&prefix) - .ok() - .is_some_and(|body| { - ["error", "errors"].iter().any(|key| { - body.get(key).is_some_and(|value| { - !value.is_null() && value != false && value != 0 && value != "" - }) - }) - }); ensure!( - !json_error, + !error_document.is_error(), "Diagnostic endpoint returned an error document" ); ensure!(!speed || bytes > 0, "Speed response contained no data"); diff --git a/src/diagnostics/error_document.rs b/src/diagnostics/error_document.rs new file mode 100644 index 0000000..d3aab03 --- /dev/null +++ b/src/diagnostics/error_document.rs @@ -0,0 +1,170 @@ +//! Recognize top-level JSON error markers without buffering response values. +//! This classifies provider error documents; it does not validate arbitrary JSON. + +#[derive(Default)] +pub(super) struct ErrorDocument { + started: bool, + finished: bool, + depth: usize, + string: bool, + escaped: bool, + expect_key: bool, + key: Option<Vec<u8>>, + error_key: bool, + value: Option<ErrorValue>, + detected: bool, +} + +enum ErrorValue { + Start, + String, + Container(u8), + Literal(Vec<u8>), +} + +impl ErrorDocument { + pub fn feed(&mut self, bytes: &[u8]) { + for &byte in bytes { + if self.finished || self.detected { + break; + } + if !self.started { + // JSON producers sometimes include a UTF-8 BOM. + if byte.is_ascii_whitespace() || matches!(byte, 0xef | 0xbb | 0xbf) { + continue; + } + self.started = true; + if byte != b'{' { + self.finished = true; + continue; + } + self.depth = 1; + self.expect_key = true; + continue; + } + self.inspect_value(byte); + if self.detected { + break; + } + if self.string { + if let Some(key) = &mut self.key { + if key.len() < 64 { + key.push(byte); + } else { + self.key = None; + } + } + if self.escaped { + self.escaped = false; + continue; + } + if byte == b'\\' { + self.escaped = true; + continue; + } + if byte == b'"' { + self.string = false; + if let Some(key) = self.key.take() { + self.error_key = + serde_json::from_slice::<String>(&key) + .ok() + .is_some_and(|key| { + key.eq_ignore_ascii_case("error") + || key.eq_ignore_ascii_case("errors") + }); + self.expect_key = false; + } + } + continue; + } + match byte { + b'"' => { + self.string = true; + if self.depth == 1 && self.expect_key { + self.key = Some(vec![b'"']); + } + } + b'{' | b'[' => self.depth = self.depth.saturating_add(1), + b'}' | b']' => { + self.depth = self.depth.saturating_sub(1); + if self.depth == 0 { + self.finished = true; + } + } + b':' if self.depth == 1 && self.error_key => { + self.value = Some(ErrorValue::Start); + self.error_key = false; + } + b',' if self.depth == 1 => { + self.expect_key = true; + self.error_key = false; + } + _ => {} + } + } + } + + fn inspect_value(&mut self, byte: u8) { + let Some(value) = self.value.take() else { + return; + }; + self.value = match value { + ErrorValue::Start if byte.is_ascii_whitespace() => Some(ErrorValue::Start), + ErrorValue::Start => match byte { + b'"' => Some(ErrorValue::String), + b'[' => Some(ErrorValue::Container(b']')), + b'{' => Some(ErrorValue::Container(b'}')), + b'n' | b'f' | b'0' | b'-' => Some(ErrorValue::Literal(vec![byte])), + _ => { + self.detected = true; + None + } + }, + ErrorValue::String => { + self.detected = byte != b'"'; + None + } + ErrorValue::Container(end) if byte.is_ascii_whitespace() => { + Some(ErrorValue::Container(end)) + } + ErrorValue::Container(end) => { + self.detected = byte != end; + None + } + ErrorValue::Literal(mut bytes) => { + if byte.is_ascii_whitespace() || matches!(byte, b',' | b'}') { + self.detected = literal_is_error(&bytes); + None + } else if bytes.len() < 64 { + bytes.push(byte); + Some(ErrorValue::Literal(bytes)) + } else { + self.detected = true; + None + } + } + }; + } + + pub fn is_error(&self) -> bool { + self.detected + || match &self.value { + Some(ErrorValue::Literal(bytes)) => literal_is_error(bytes), + // An incomplete value of a declared error field is not a successful download. + Some(_) => true, + None => false, + } + } + + pub fn confirmed_error(&self) -> bool { + self.detected + } +} + +fn literal_is_error(bytes: &[u8]) -> bool { + match serde_json::from_slice::<serde_json::Value>(bytes) { + Ok(serde_json::Value::Null | serde_json::Value::Bool(false)) => false, + Ok(serde_json::Value::Number(number)) => number.as_f64() != Some(0.0), + _ => true, + } +} |