fix: block ALL LS follow-up requests across connections

Move the in-flight blocking check to the top of the LLM request flow, BEFORE request modification. This catches follow-ups on ALL connections (the LS opens multiple parallel TLS connections). Only the very first modified request reaches Google — all others get fake STOP responses. Previously, each new connection independently allowed one request through before blocking, letting 4-5 requests leak per turn.
2026-02-16 00:57:33 -06:00
parent a8f3c8915f
commit 3fdd0368a0
23 changed files with 992 additions and 568 deletions
--- a/src/snapshot.rs
+++ b/src/snapshot.rs
@@ -10,16 +10,44 @@ use std::io::{self, Read};
 // ── Domain metadata ──────────────────────────────────────────────────────────

 const DOMAIN_INFO: &[(&str, &str, &str)] = &[
-    ("antigravity-unleash.goog", "Feature Flags", "Unleash SDK — controls A/B tests and feature rollouts"),
-    ("daily-cloudcode-pa.googleapis.com", "LLM API (gRPC)", "Primary Gemini/Claude API endpoint"),
-    ("cloudcode-pa.googleapis.com", "LLM API (gRPC)", "Production Gemini/Claude API endpoint"),
-    ("api.anthropic.com", "Claude API", "Direct Anthropic API calls"),
-    ("lh3.googleusercontent.com", "Profile Picture", "User avatar"),
+    (
+        "antigravity-unleash.goog",
+        "Feature Flags",
+        "Unleash SDK — controls A/B tests and feature rollouts",
+    ),
+    (
+        "daily-cloudcode-pa.googleapis.com",
+        "LLM API (gRPC)",
+        "Primary Gemini/Claude API endpoint",
+    ),
+    (
+        "cloudcode-pa.googleapis.com",
+        "LLM API (gRPC)",
+        "Production Gemini/Claude API endpoint",
+    ),
+    (
+        "api.anthropic.com",
+        "Claude API",
+        "Direct Anthropic API calls",
+    ),
+    (
+        "lh3.googleusercontent.com",
+        "Profile Picture",
+        "User avatar",
+    ),
    ("play.googleapis.com", "Telemetry", "Google Play telemetry"),
-    ("firebaseinstallations.googleapis.com", "Firebase", "Installation tracking"),
+    (
+        "firebaseinstallations.googleapis.com",
+        "Firebase",
+        "Installation tracking",
+    ),
    ("oauth2.googleapis.com", "OAuth", "Token refresh/exchange"),
    ("speech.googleapis.com", "Speech", "Voice input processing"),
-    ("modelarmor.googleapis.com", "Safety", "Content safety/filtering"),
+    (
+        "modelarmor.googleapis.com",
+        "Safety",
+        "Content safety/filtering",
+    ),
 ];

 fn domain_label(domain: &str) -> (&str, &str) {
@@ -57,8 +85,8 @@ struct HttpExchange {

 #[derive(Debug, Clone, Copy, PartialEq)]
 enum Direction {
-    Outgoing,  // LS → upstream
-    Incoming,  // external → LS (our curl calls)
+    Outgoing, // LS → upstream
+    Incoming, // external → LS (our curl calls)
 }

 #[derive(Default)]
@@ -101,10 +129,12 @@ impl Snapshot {

            // LS process logs
            if (line.starts_with('I') || line.starts_with('W') || line.starts_with('E'))
-                && line.len() > 4 && line.chars().nth(1).is_some_and(|c| c.is_ascii_digit()) {
-                    snap.ls_logs.push(line.to_string());
-                    continue;
-                }
+                && line.len() > 4
+                && line.chars().nth(1).is_some_and(|c| c.is_ascii_digit())
+            {
+                snap.ls_logs.push(line.to_string());
+                continue;
+            }
            if line.contains("maxprocs:") {
                snap.ls_logs.push(line.to_string());
                continue;
@@ -128,8 +158,15 @@ impl Snapshot {
            if let Some((key, val)) = extract_header(line, "Transport encoding header") {
                if key == ":method" {
                    // Finalize previous exchange
-                    if current_pseudo.contains_key(":path") || current_pseudo.contains_key(":method") {
-                        snap.finalize_exchange(&current_pseudo, &current_headers, current_direction, current_stream.clone());
+                    if current_pseudo.contains_key(":path")
+                        || current_pseudo.contains_key(":method")
+                    {
+                        snap.finalize_exchange(
+                            &current_pseudo,
+                            &current_headers,
+                            current_direction,
+                            current_stream.clone(),
+                        );
                    }
                    current_headers.clear();
                    current_pseudo.clear();
@@ -147,8 +184,15 @@ impl Snapshot {
            // Incoming / server-received headers
            if let Some((key, val)) = extract_header(line, "decoded hpack field header field") {
                if key == ":authority" && !line.contains("server read frame") {
-                    if current_pseudo.contains_key(":path") || current_pseudo.contains_key(":method") {
-                        snap.finalize_exchange(&current_pseudo, &current_headers, current_direction, current_stream.clone());
+                    if current_pseudo.contains_key(":path")
+                        || current_pseudo.contains_key(":method")
+                    {
+                        snap.finalize_exchange(
+                            &current_pseudo,
+                            &current_headers,
+                            current_direction,
+                            current_stream.clone(),
+                        );
                    }
                    current_headers.clear();
                    current_pseudo.clear();
@@ -167,8 +211,15 @@ impl Snapshot {
            if line.contains("wrote HEADERS") {
                if let Some(stream) = extract_stream_id(line) {
                    current_stream = Some(stream.clone());
-                    if current_pseudo.contains_key(":path") || current_pseudo.contains_key(":method") {
-                        let ex = snap.finalize_exchange(&current_pseudo, &current_headers, current_direction, Some(stream));
+                    if current_pseudo.contains_key(":path")
+                        || current_pseudo.contains_key(":method")
+                    {
+                        let ex = snap.finalize_exchange(
+                            &current_pseudo,
+                            &current_headers,
+                            current_direction,
+                            Some(stream),
+                        );
                        if ex.is_some() {
                            current_headers.clear();
                            current_pseudo.clear();
@@ -179,10 +230,13 @@ impl Snapshot {
            }

            // DATA frames
-            if (line.contains("wrote DATA") || line.contains("read DATA") || line.contains("server read frame DATA"))
+            if (line.contains("wrote DATA")
+                || line.contains("read DATA")
+                || line.contains("server read frame DATA"))
                && line.contains("data=\"")
            {
-                let is_outgoing = line.contains("wrote DATA") || line.contains("server read frame DATA");
+                let is_outgoing =
+                    line.contains("wrote DATA") || line.contains("server read frame DATA");
                if let Some(stream) = extract_stream_id(line) {
                    if let Some(data_str) = extract_data(line) {
                        let raw = decode_go_escaped(&data_str);
@@ -203,7 +257,12 @@ impl Snapshot {

        // Finalize remaining
        if current_pseudo.contains_key(":path") || current_pseudo.contains_key(":method") {
-            snap.finalize_exchange(&current_pseudo, &current_headers, current_direction, current_stream);
+            snap.finalize_exchange(
+                &current_pseudo,
+                &current_headers,
+                current_direction,
+                current_stream,
+            );
        }

        snap
@@ -226,7 +285,11 @@ impl Snapshot {

        self.exchanges.push(HttpExchange {
            authority,
-            method: if method.is_empty() { "GET".into() } else { method },
+            method: if method.is_empty() {
+                "GET".into()
+            } else {
+                method
+            },
            path,
            headers: headers.to_vec(),
            body: Vec::new(),
@@ -245,7 +308,9 @@ impl Snapshot {
        let sep = "═".repeat(70);
        let sep_thin = "─".repeat(60);
        out.push_str(&format!("\n{BOLD}{CYAN}{sep}{NC}\n"));
-        out.push_str(&format!("{BOLD}{CYAN}  STANDALONE LS TRAFFIC SNAPSHOT{NC}\n"));
+        out.push_str(&format!(
+            "{BOLD}{CYAN}  STANDALONE LS TRAFFIC SNAPSHOT{NC}\n"
+        ));
        out.push_str(&format!("{BOLD}{CYAN}{sep}{NC}\n\n"));

        // LS Logs
@@ -265,7 +330,9 @@ impl Snapshot {
            for target in &self.connections {
                let domain = target.split(':').next().unwrap_or(target);
                let (label, desc) = domain_label(domain);
-                out.push_str(&format!("  {GREEN}→{NC} {BOLD}{target}{NC}  {DIM}({label}){NC}\n"));
+                out.push_str(&format!(
+                    "  {GREEN}→{NC} {BOLD}{target}{NC}  {DIM}({label}){NC}\n"
+                ));
                if !desc.is_empty() {
                    out.push_str(&format!("    {DIM}{desc}{NC}\n"));
                }
@@ -276,7 +343,10 @@ impl Snapshot {
        // Group by domain
        let mut by_domain: Vec<(&str, Vec<&HttpExchange>)> = Vec::new();
        for ex in &self.exchanges {
-            if let Some(entry) = by_domain.iter_mut().find(|(d, _)| *d == ex.authority.as_str()) {
+            if let Some(entry) = by_domain
+                .iter_mut()
+                .find(|(d, _)| *d == ex.authority.as_str())
+            {
                entry.1.push(ex);
            } else {
                by_domain.push((&ex.authority, vec![ex]));
@@ -293,12 +363,17 @@ impl Snapshot {
            let color = if label.contains("API") { YELLOW } else { CYAN };

            out.push_str(&format!("\n{BOLD}{sep}{NC}\n"));
-            out.push_str(&format!("{BOLD}{color}  {domain}{NC}  {DIM}— {label}{NC}\n"));
+            out.push_str(&format!(
+                "{BOLD}{color}  {domain}{NC}  {DIM}— {label}{NC}\n"
+            ));
            out.push_str(&format!("{BOLD}{sep}{NC}\n"));

            for ex in exchanges {
                let method_color = if ex.method == "GET" { GREEN } else { YELLOW };
-                out.push_str(&format!("\n  {BOLD}→ {method_color}{}{NC} {}\n", ex.method, ex.path));
+                out.push_str(&format!(
+                    "\n  {BOLD}→ {method_color}{}{NC} {}\n",
+                    ex.method, ex.path
+                ));

                // Interesting headers
                for (key, val) in &ex.headers {
@@ -342,7 +417,10 @@ fn render_body(data: &[u8], total_len: usize) -> String {
            out.push_str(&format!("    {BOLD}Body ({len} bytes, JSON):{NC}\n"));
            for (i, line) in pretty.lines().enumerate() {
                if i >= 40 {
-                    out.push_str(&format!("      {DIM}... ({} more lines){NC}\n", pretty.lines().count() - 40));
+                    out.push_str(&format!(
+                        "      {DIM}... ({} more lines){NC}\n",
+                        pretty.lines().count() - 40
+                    ));
                    break;
                }
                out.push_str(&format!("      {GREEN}{line}{NC}\n"));
@@ -357,10 +435,16 @@ fn render_body(data: &[u8], total_len: usize) -> String {
            if let Ok(text) = std::str::from_utf8(&decompressed) {
                if let Ok(val) = serde_json::from_str::<serde_json::Value>(text) {
                    let pretty = serde_json::to_string_pretty(&val).unwrap_or_default();
-                    out.push_str(&format!("    {BOLD}Body ({len} bytes gzip → {} bytes, JSON):{NC}\n", decompressed.len()));
+                    out.push_str(&format!(
+                        "    {BOLD}Body ({len} bytes gzip → {} bytes, JSON):{NC}\n",
+                        decompressed.len()
+                    ));
                    for (i, line) in pretty.lines().enumerate() {
                        if i >= 50 {
-                            out.push_str(&format!("      {DIM}... ({} more lines){NC}\n", pretty.lines().count() - 50));
+                            out.push_str(&format!(
+                                "      {DIM}... ({} more lines){NC}\n",
+                                pretty.lines().count() - 50
+                            ));
                            break;
                        }
                        out.push_str(&format!("      {GREEN}{line}{NC}\n"));
@@ -368,14 +452,20 @@ fn render_body(data: &[u8], total_len: usize) -> String {
                    return out;
                }
                // Plain text
-                out.push_str(&format!("    {BOLD}Body ({len} bytes gzip → {} bytes, text):{NC}\n", decompressed.len()));
+                out.push_str(&format!(
+                    "    {BOLD}Body ({len} bytes gzip → {} bytes, text):{NC}\n",
+                    decompressed.len()
+                ));
                for line in text.lines().take(20) {
                    out.push_str(&format!("      {line}\n"));
                }
                return out;
            }
            // Binary gzip
-            out.push_str(&format!("    {BOLD}Body ({len} bytes gzip → {} bytes, binary):{NC}\n", decompressed.len()));
+            out.push_str(&format!(
+                "    {BOLD}Body ({len} bytes gzip → {} bytes, binary):{NC}\n",
+                decompressed.len()
+            ));
            let strings = extract_strings(&decompressed);
            for s in strings.iter().take(15) {
                out.push_str(&format!("      {MAGENTA}{s}{NC}\n"));
@@ -393,7 +483,11 @@ fn render_body(data: &[u8], total_len: usize) -> String {
    // Protobuf / binary with string extraction
    let strings = extract_strings(data);
    if !strings.is_empty() {
-        let kind = if !data.is_empty() && matches!(data[0], 0x08 | 0x0a | 0x10 | 0x12 | 0x18 | 0x1a | 0x20 | 0x22) {
+        let kind = if !data.is_empty()
+            && matches!(
+                data[0],
+                0x08 | 0x0a | 0x10 | 0x12 | 0x18 | 0x1a | 0x20 | 0x22
+            ) {
            "protobuf"
        } else {
            "binary"
@@ -448,7 +542,9 @@ fn extract_header(line: &str, pattern: &str) -> Option<(String, String)> {
 fn extract_stream_id(line: &str) -> Option<String> {
    let pos = line.find("stream=")?;
    let rest = &line[pos + 7..];
-    let end = rest.find(|c: char| !c.is_ascii_digit()).unwrap_or(rest.len());
+    let end = rest
+        .find(|c: char| !c.is_ascii_digit())
+        .unwrap_or(rest.len());
    Some(rest[..end].to_string())
 }

@@ -470,7 +566,9 @@ fn extract_data(line: &str) -> Option<String> {
 fn extract_data_len(line: &str) -> Option<usize> {
    let pos = line.find("len=")?;
    let rest = &line[pos + 4..];
-    let end = rest.find(|c: char| !c.is_ascii_digit()).unwrap_or(rest.len());
+    let end = rest
+        .find(|c: char| !c.is_ascii_digit())
+        .unwrap_or(rest.len());
    rest[..end].parse().ok()
 }

@@ -482,17 +580,40 @@ fn decode_go_escaped(s: &str) -> Vec<u8> {
        if bytes[i] == b'\\' && i + 1 < bytes.len() {
            match bytes[i + 1] {
                b'x' if i + 3 < bytes.len() => {
-                    if let Ok(b) = u8::from_str_radix(std::str::from_utf8(&bytes[i + 2..i + 4]).unwrap_or(""), 16) {
+                    if let Ok(b) = u8::from_str_radix(
+                        std::str::from_utf8(&bytes[i + 2..i + 4]).unwrap_or(""),
+                        16,
+                    ) {
                        result.push(b);
                        i += 4;
                        continue;
                    }
                }
-                b'n' => { result.push(b'\n'); i += 2; continue; }
-                b'r' => { result.push(b'\r'); i += 2; continue; }
-                b't' => { result.push(b'\t'); i += 2; continue; }
-                b'\\' => { result.push(b'\\'); i += 2; continue; }
-                b'"' => { result.push(b'"'); i += 2; continue; }
+                b'n' => {
+                    result.push(b'\n');
+                    i += 2;
+                    continue;
+                }
+                b'r' => {
+                    result.push(b'\r');
+                    i += 2;
+                    continue;
+                }
+                b't' => {
+                    result.push(b'\t');
+                    i += 2;
+                    continue;
+                }
+                b'\\' => {
+                    result.push(b'\\');
+                    i += 2;
+                    continue;
+                }
+                b'"' => {
+                    result.push(b'"');
+                    i += 2;
+                    continue;
+                }
                _ => {}
            }
        }
@@ -562,7 +683,10 @@ pub fn run_cli() {
        })
    } else {
        let mut buf = String::new();
-        io::stdin().lock().read_to_string(&mut buf).expect("Failed to read stdin");
+        io::stdin()
+            .lock()
+            .read_to_string(&mut buf)
+            .expect("Failed to read stdin");
        buf
    };