Reject streamed diagnostic error documents and record VM qualification

1 files changed, 20 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");