# MITM Proxy

## Overview

The built-in MITM proxy intercepts all traffic between the standalone LS and Google's API. It decrypts TLS, parses SSE responses, captures real token usage, and modifies requests to inject tools, parameters, and images.

```mermaid
sequenceDiagram
    participant LS as Standalone LS
    participant IPT as iptables
    participant MITM as MITM Proxy :8742
    participant Store as MitmStore
    participant G as Google API

    LS->>IPT: HTTPS :443
    IPT->>MITM: REDIRECT (UID-scoped)
    MITM->>MITM: TLS terminate (dynamic cert)
    MITM->>Store: Match request by cascade_id
    Store-->>MITM: RequestContext (tools, params, image)
    MITM->>MITM: modify_request()
    MITM->>G: Forward modified request
    G-->>MITM: SSE stream
    MITM->>MITM: Parse SSE, extract usage
    MITM->>Store: Dispatch events (TextDelta, Usage, etc.)
    MITM-->>LS: Forward original response
```

---

## Components

```mermaid
graph TD
    subgraph "MITM Module"
        proxy["proxy.rs<br/>TLS termination<br/>SNI-based routing"]
        h2["h2_handler.rs<br/>HTTP/2 frame handling"]
        intercept["intercept.rs<br/>SSE parser<br/>Usage extraction"]
        modify["modify.rs<br/>Request injection<br/>(tools, params, images)"]
        store["store.rs<br/>MitmStore<br/>Event channels"]
        proto["proto.rs<br/>Protobuf codec"]
        ca["ca.rs<br/>CA + dynamic certs"]
    end

    proxy --> h2
    h2 --> intercept
    h2 --> modify
    modify --> store
    intercept --> store
    proxy --> ca
    modify --> proto

    style store fill:#dc2626,color:#fff
    style proxy fill:#ea580c,color:#fff
```

| File            | Purpose                                                                                       |
| --------------- | --------------------------------------------------------------------------------------------- |
| `proxy.rs`      | Accepts iptables-redirected connections, terminates TLS via SNI, manages connection lifecycle |
| `h2_handler.rs` | HTTP/2 frame-level handling for CONNECT-style proxying                                        |
| `intercept.rs`  | Parses Google's SSE `data:` lines, extracts `usageMetadata`, detects `finishReason`           |
| `modify.rs`     | Injects tools, generation params, images, tool results, Google Search grounding into requests |
| `store.rs`      | Central state — `RequestContext` registry, event channels (`MitmEvent`), usage accumulation   |
| `proto.rs`      | Protobuf encode/decode for intercepted request/response bodies                                |
| `ca.rs`         | Generates CA certificate and per-domain leaf certs for TLS termination                        |

---

## Request Modification

When the MITM proxy intercepts an outgoing request from the LS, it applies modifications from the `RequestContext` stored by the API handler:

```mermaid
flowchart TD
    A["Original LS Request"] --> B{"Has tools?"}
    B -- Yes --> C["Inject tool definitions<br/>+ toolConfig"]
    B -- No --> D{"Has generation params?"}
    C --> D
    D -- Yes --> E["Inject temperature, top_p,<br/>max_output_tokens, stop_sequences,<br/>frequency/presence_penalty"]
    D -- No --> F{"Has image?"}
    E --> F
    F -- Yes --> G["Inject inlineData<br/>(base64) into user parts"]
    F -- No --> H{"Has tool results?"}
    G --> H
    H -- Yes --> I["Inject functionResponse<br/>parts"]
    H -- No --> J{"Google Search?"}
    I --> J
    J -- Yes --> K["Enable Google Search<br/>grounding tool"]
    J -- No --> L["Replace user text<br/>with real input"]
    K --> L
    L --> M["Update Content-Length"]
    M --> N["Forward to Google"]

    style A fill:#2563eb,color:#fff
    style N fill:#059669,color:#fff
```

---

## SSE Response Format

Google's API returns SSE events:

```
data: {"response": {"candidates": [{"content": {"role": "model", "parts": [{"text": "..."}]}}],
       "usageMetadata": {"promptTokenCount": 1514, "candidatesTokenCount": 25,
                         "totalTokenCount": 1539, "thoughtsTokenCount": 52},
       "modelVersion": "gemini-3-flash"}, "traceId": "...", "metadata": {}}
```

The last event includes `"finishReason": "STOP"` in the candidate.

---

## MitmEvent Channel

Events dispatched through `tokio::sync::mpsc` channels from MITM → API handlers:

| Event                   | Source         | Data                                          |
| ----------------------- | -------------- | --------------------------------------------- |
| `TextDelta(String)`     | `intercept.rs` | Incremental text from model                   |
| `ThinkingDelta(String)` | `intercept.rs` | Thinking/reasoning text                       |
| `Usage(ApiUsage)`       | `intercept.rs` | Token counts (input, output, thinking, cache) |
| `FunctionCall(Vec)`     | `intercept.rs` | Tool calls from model                         |
| `Grounding(Value)`      | `intercept.rs` | Google Search grounding metadata              |
| `ResponseComplete`      | `intercept.rs` | Stream finished                               |
| `UpstreamError(Value)`  | `intercept.rs` | Google API error (400, 429, 500)              |

---

## Setup

### UID-Scoped iptables (Classic Mode)

```bash
# One-time setup — creates zerogravity-ls user + iptables rule
sudo ./scripts/mitm-redirect.sh install

# Run proxy (standalone LS + MITM both enabled by default)
RUST_LOG=info ./target/release/zerogravity

# Check intercepted usage
curl -s http://localhost:8741/v1/usage | jq .

# Cleanup
sudo ./scripts/mitm-redirect.sh uninstall
```

### Headless Mode

No iptables or sudo needed. The LS connects through `HTTPS_PROXY` instead:

```bash
RUST_LOG=info ./target/release/zerogravity --headless
```

---

## Intercepted Endpoints

| Endpoint                    | Type     | Content                   |
| --------------------------- | -------- | ------------------------- |
| `streamGenerateContent`     | SSE/JSON | LLM responses ✅ (parsed) |
| `fetchUserInfo`             | Protobuf | User info                 |
| `loadCodeAssist`            | Protobuf | Extension config          |
| `fetchAvailableModels`      | Protobuf | Model catalog             |
| `recordCodeAssistMetrics`   | Protobuf | Telemetry (ignored)       |
| `recordTrajectoryAnalytics` | Protobuf | Telemetry (ignored)       |
