b3nw/zerogravity
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: MITM interception for standalone LS with UID isolation

Browse Source 
- Spawn standalone LS as dedicated 'antigravity-ls' user via sudo
- UID-scoped iptables redirect (port 443 → MITM proxy) via mitm-redirect.sh
- Combined CA bundle (system CAs + MITM CA) for Go TLS trust
- Transparent TLS interception with chunked response detection
- Google SSE parser for streamGenerateContent usage extraction
- Timeouts on all MITM operations (TLS handshake, upstream, idle)
- Forward response data immediately (no buffering)
- Per-model token usage capture (input, output, thinking)
- Update docs and known issues to reflect resolved TLS blocker
This commit is contained in:
Nikketryhard
2026-02-14 17:50:12 -06:00
parent 6842bfeaa5
commit d4de436856
10 changed files with 1156 additions and 478 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
src/backend.rs
View File

@@ -83,6 +83,34 @@ impl Backend {
})
}
/// Create a Backend with known connection details (for standalone LS).
///
/// Skips auto-discovery — the caller provides the port, CSRF, and OAuth token.
pub fn new_with_config(
port: u16,
csrf: String,
oauth_token: String,
) -> Result<Self, String> {
let inner = BackendInner {
pid: "standalone".to_string(),
csrf,
https_port: port.to_string(),
oauth_token,
};
let client = wreq::Client::builder()
.emulation(wreq_util::Emulation::Chrome142)
.cert_verification(false)
.verify_hostname(false)
.build()
.map_err(|e| format!("wreq client build failed: {e}"))?;
Ok(Self {
inner: RwLock::new(inner),
client,
})
}
/// Re-discover language server connection details.
/// Runs blocking I/O on a spawn_blocking thread to avoid starving tokio.
pub async fn refresh(&self) -> Result<(), String> {

99
src/main.rs
View File

@@ -11,6 +11,7 @@ mod mitm;
mod proto;
mod quota;
mod session;
mod standalone;
mod warmup;
use api::AppState;
@@ -44,6 +45,10 @@ struct Cli {
/// MITM proxy port (default: 8742, matches wrapper script)
#[arg(long, default_value_t = 8742)]
mitm_port: u16,
/// Use a standalone LS (does not touch the real LS)
#[arg(long)]
standalone: bool,
}
#[tokio::main]
@@ -85,12 +90,83 @@ async fn main() {
}
};
// ── Step 2: Backend discovery ─────────────────────────────────────────────
let backend = Arc::new(match Backend::new() {
Ok(b) => b,
Err(e) => {
eprintln!("Fatal: {e}");
std::process::exit(1);
// ── Step 2: Backend discovery (or standalone LS spawn) ─────────────────────
let standalone_ls = if cli.standalone {
// Standalone mode: discover main LS config, spawn our own
let main_config = match standalone::discover_main_ls_config() {
Ok(c) => c,
Err(e) => {
eprintln!("Fatal: {e}");
std::process::exit(1);
}
};
// Build MITM config if MITM is enabled
let mitm_cfg = if !cli.no_mitm {
let ca_path = dirs_data_dir()
.join("mitm-ca.pem")
.to_string_lossy()
.to_string();
Some(standalone::StandaloneMitmConfig {
proxy_addr: format!("http://127.0.0.1:{}", cli.mitm_port),
ca_cert_path: ca_path,
})
} else {
None
};
let ls = match standalone::StandaloneLS::spawn(&main_config, mitm_cfg.as_ref()) {
Ok(ls) => ls,
Err(e) => {
eprintln!("Fatal: failed to spawn standalone LS: {e}");
std::process::exit(1);
}
};
// Wait for it to be ready
let rt_ls_port = ls.port;
let rt_ls_csrf = ls.csrf.clone();
tokio::task::block_in_place(|| {
tokio::runtime::Handle::current().block_on(async {
if let Err(e) = ls.wait_ready(10).await {
eprintln!("Fatal: {e}");
std::process::exit(1);
}
});
});
info!(port = rt_ls_port, "Standalone LS ready");
Some((ls, rt_ls_port, rt_ls_csrf))
} else {
None
};
let backend = Arc::new(if let Some((_, port, ref csrf)) = standalone_ls {
// Build backend pointing at standalone LS
let oauth = std::env::var("ANTIGRAVITY_OAUTH_TOKEN")
.ok()
.filter(|s| !s.is_empty())
.or_else(|| {
let home = std::env::var("HOME").unwrap_or_default();
let path = format!("{home}/.config/antigravity-proxy-token");
std::fs::read_to_string(&path)
.ok()
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
})
.unwrap_or_default();
match Backend::new_with_config(port, csrf.clone(), oauth) {
Ok(b) => b,
Err(e) => {
eprintln!("Fatal: {e}");
std::process::exit(1);
}
}
} else {
// Normal mode: discover existing LS
match Backend::new() {
Ok(b) => b,
Err(e) => {
eprintln!("Fatal: {e}");
std::process::exit(1);
}
}
});
@@ -151,8 +227,15 @@ async fn main() {
});
// Periodic backend refresh — keeps LS connection details fresh
// (skip in standalone mode — the port is fixed and discover() would overwrite it)
let is_standalone = cli.standalone;
let refresh_backend = Arc::clone(&state.backend);
let refresh_handle = tokio::spawn(async move {
if is_standalone {
// In standalone mode, the backend config is fixed — no refresh needed
std::future::pending::<()>().await;
return;
}
loop {
tokio::time::sleep(tokio::time::Duration::from_secs(60)).await;
if let Err(e) = refresh_backend.refresh().await {
@@ -178,6 +261,10 @@ async fn main() {
if let Some(h) = mitm_handle {
h.abort();
}
// Kill standalone LS if we spawned one
if let Some((mut ls, _, _)) = standalone_ls {
ls.kill();
}
// Remove stale MITM port file
let _ = std::fs::remove_file(dirs_data_dir().join("mitm-port"));
info!("Server shutdown complete");

60
src/mitm/intercept.rs
View File

@@ -56,9 +56,11 @@ pub struct StreamingAccumulator {
pub output_tokens: u64,
pub cache_creation_input_tokens: u64,
pub cache_read_input_tokens: u64,
pub thinking_tokens: u64,
pub model: Option<String>,
pub stop_reason: Option<String>,
pub is_complete: bool,
pub api_provider: Option<String>,
}
impl StreamingAccumulator {
@@ -66,13 +68,46 @@ impl StreamingAccumulator {
Self::default()
}
/// Process a single SSE event.
/// Process a single SSE event.
pub fn process_event(&mut self, event: &Value) {
// ── Google format: {"response": {"usageMetadata": {...}, "modelVersion": "..."}} ──
if let Some(response) = event.get("response") {
// Extract usage metadata (each event has cumulative counts)
if let Some(usage) = response.get("usageMetadata") {
self.input_tokens = usage["promptTokenCount"].as_u64().unwrap_or(self.input_tokens);
self.output_tokens = usage["candidatesTokenCount"].as_u64().unwrap_or(self.output_tokens);
self.thinking_tokens = usage["thoughtsTokenCount"].as_u64().unwrap_or(self.thinking_tokens);
}
if let Some(model) = response["modelVersion"].as_str() {
self.model = Some(model.to_string());
}
// Check for completion in candidates
if let Some(candidates) = response.get("candidates").and_then(|c| c.as_array()) {
for candidate in candidates {
if let Some(reason) = candidate["finishReason"].as_str() {
self.stop_reason = Some(reason.to_string());
if reason == "STOP" {
self.is_complete = true;
}
}
}
}
self.api_provider = Some("google".to_string());
trace!(
input = self.input_tokens,
output = self.output_tokens,
thinking = self.thinking_tokens,
complete = self.is_complete,
"SSE Google: usage update"
);
return;
}
// ── Anthropic format: {"type": "message_start"|"message_delta"|"message_stop"} ──
let event_type = event["type"].as_str().unwrap_or("");
match event_type {
"message_start" => {
// message_start contains the initial usage (input tokens + cache)
if let Some(usage) = event.get("message").and_then(|m| m.get("usage")) {
self.input_tokens = usage["input_tokens"].as_u64().unwrap_or(0);
self.cache_creation_input_tokens = usage["cache_creation_input_tokens"].as_u64().unwrap_or(0);
@@ -81,36 +116,27 @@ impl StreamingAccumulator {
if let Some(model) = event.get("message").and_then(|m| m["model"].as_str()) {
self.model = Some(model.to_string());
}
trace!(
input = self.input_tokens,
cache_read = self.cache_read_input_tokens,
cache_create = self.cache_creation_input_tokens,
"SSE message_start: captured input usage"
);
self.api_provider = Some("anthropic".to_string());
trace!(input = self.input_tokens, "SSE Anthropic: message_start");
}
"message_delta" => {
// message_delta contains the output usage
if let Some(usage) = event.get("usage") {
self.output_tokens = usage["output_tokens"].as_u64().unwrap_or(self.output_tokens);
}
if let Some(reason) = event["delta"]["stop_reason"].as_str() {
self.stop_reason = Some(reason.to_string());
}
trace!(output = self.output_tokens, "SSE message_delta: updated output tokens");
}
"message_stop" => {
self.is_complete = true;
debug!(
input = self.input_tokens,
output = self.output_tokens,
cache_read = self.cache_read_input_tokens,
model = ?self.model,
"SSE message_stop: stream complete"
"SSE Anthropic: stream complete"
);
}
"content_block_start" | "content_block_delta" | "content_block_stop" | "ping" => {
// Content events — no usage data, just pass through
}
"content_block_start" | "content_block_delta" | "content_block_stop" | "ping" => {}
_ => {
trace!(event_type, "SSE: unknown event type");
}
@@ -124,11 +150,11 @@ impl StreamingAccumulator {
output_tokens: self.output_tokens,
cache_creation_input_tokens: self.cache_creation_input_tokens,
cache_read_input_tokens: self.cache_read_input_tokens,
thinking_output_tokens: 0,
thinking_output_tokens: self.thinking_tokens,
response_output_tokens: 0,
model: self.model,
stop_reason: self.stop_reason,
api_provider: Some("anthropic".to_string()),
api_provider: self.api_provider.unwrap_or_else(|| "unknown".to_string()).into(),
grpc_method: None,
captured_at: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)

285
src/mitm/proxy.rs
View File

@@ -85,7 +85,7 @@ pub async fn run(
let store = store.clone();
tokio::spawn(async move {
if let Err(e) = handle_connection(stream, ca, store, modify_requests).await {
debug!(error = %e, "MITM connection error");
warn!(error = %e, "MITM connection error");
}
});
}
@@ -310,18 +310,30 @@ async fn handle_intercepted(
let acceptor = TlsAcceptor::from(server_config);
// Perform TLS handshake with the client (LS)
let tls_stream = acceptor
.accept(stream)
.await
.map_err(|e| format!("TLS handshake with client failed for {domain}: {e}"))?;
// Perform TLS handshake with the client (LS) — 10s timeout
let tls_stream = match tokio::time::timeout(
std::time::Duration::from_secs(10),
acceptor.accept(stream),
)
.await
{
Ok(Ok(s)) => s,
Ok(Err(e)) => {
warn!(domain, error = %e, "MITM: TLS handshake FAILED (client rejected cert?)");
return Err(format!("TLS handshake with client failed for {domain}: {e}"));
}
Err(_) => {
warn!(domain, "MITM: TLS handshake TIMED OUT after 10s");
return Err(format!("TLS handshake timed out for {domain}"));
}
};
// Check negotiated ALPN protocol
let alpn = tls_stream.get_ref().1
.alpn_protocol()
.map(|p| String::from_utf8_lossy(p).to_string());
debug!(domain, alpn = ?alpn, "MITM: TLS handshake successful");
info!(domain, alpn = ?alpn, "MITM: TLS handshake successful ✓");
match alpn.as_deref() {
Some("h2") => {
@@ -336,7 +348,7 @@ async fn handle_intercepted(
}
_ => {
// HTTP/1.1 or no ALPN — use the existing handler
debug!(domain, "MITM: routing to HTTP/1.1 handler");
info!(domain, "MITM: routing to HTTP/1.1 handler");
handle_http_over_tls(tls_stream, domain, store, modify_requests).await
}
}
@@ -382,16 +394,35 @@ async fn handle_http_over_tls(
// Try to resolve the real IP, bypassing /etc/hosts
let addr = resolve_upstream(domain).await;
info!(domain, addr = %addr, "MITM: connecting upstream");
let tcp = match tokio::time::timeout(
std::time::Duration::from_secs(15),
TcpStream::connect(&addr),
)
.await
{
Ok(Ok(s)) => s,
Ok(Err(e)) => return Err(format!("Connect to upstream {domain} ({addr}): {e}")),
Err(_) => return Err(format!("Connect to upstream {domain} ({addr}): timed out")),
};
let tcp = TcpStream::connect(addr)
.await
.map_err(|e| format!("Connect to upstream {domain}: {e}"))?;
let server_name = rustls::pki_types::ServerName::try_from(domain.to_string())
.map_err(|e| format!("Invalid server name: {e}"))?;
connector
.connect(server_name, tcp)
.await
.map_err(|e| format!("TLS connect to upstream {domain}: {e}"))
match tokio::time::timeout(
std::time::Duration::from_secs(15),
connector.connect(server_name, tcp),
)
.await
{
Ok(Ok(s)) => {
info!(domain, "MITM: upstream TLS connected ✓");
Ok(s)
}
Ok(Err(e)) => Err(format!("TLS connect to upstream {domain}: {e}")),
Err(_) => Err(format!("TLS connect to upstream {domain}: timed out")),
}
}
/// Resolve upstream IP bypassing /etc/hosts.
@@ -428,8 +459,37 @@ async fn handle_http_over_tls(
// ── Read the HTTP request from the client ─────────────────────────
let mut request_buf = Vec::with_capacity(1024 * 64);
// 60s timeout on initial read (LS may open connection without sending immediately)
const IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60);
loop {
let n = match client.read(&mut tmp).await {
let read_result = if request_buf.is_empty() {
// First read — apply idle timeout
match tokio::time::timeout(IDLE_TIMEOUT, client.read(&mut tmp)).await {
Ok(r) => r,
Err(_) => {
// Idle timeout — connection pool warmup, no data sent
debug!(domain, "MITM: client idle timeout (60s), closing");
return Ok(());
}
}
} else {
// Subsequent reads — wait up to 30s for rest of request
match tokio::time::timeout(
std::time::Duration::from_secs(30),
client.read(&mut tmp),
)
.await
{
Ok(r) => r,
Err(_) => {
warn!(domain, "MITM: partial request read timed out");
return Err("Partial request read timed out".into());
}
}
};
let n = match read_result {
Ok(0) => return Ok(()), // Client closed connection cleanly
Ok(n) => n,
Err(e) => {
@@ -461,12 +521,25 @@ async fn handle_http_over_tls(
None
};
debug!(
// Extract request method and path for logging
let req_path = {
let mut headers = [httparse::EMPTY_HEADER; 64];
let mut req = httparse::Request::new(&mut headers);
match req.parse(&request_buf) {
Ok(httparse::Status::Complete(_)) => {
format!("{} {}", req.method.unwrap_or("?"), req.path.unwrap_or("?"))
}
_ => "?".to_string(),
}
};
info!(
domain,
req_path = %req_path,
content_length,
streaming = is_streaming_request,
cascade = ?cascade_hint,
"MITM: forwarding request to upstream"
"MITM: forwarding request"
);
// ── Ensure upstream connection is alive ──────────────────────────────
@@ -492,118 +565,139 @@ async fn handle_http_over_tls(
let conn = upstream.as_mut().unwrap();
// ── Stream response back to client ──────────────────────────────────
// ALWAYS forward data to client immediately (no buffering).
// Buffer body on the side for usage parsing.
let mut streaming_acc = StreamingAccumulator::new();
let mut is_streaming_response = false;
let mut headers_parsed = false;
// Only buffer response body for non-streaming (for usage parsing)
let mut non_streaming_buf: Option<Vec<u8>> = None;
// Track if upstream connection is still usable after this response
let mut upstream_ok = true;
// Per-request timeout: 5 minutes (covers large context API calls)
const READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300);
let mut response_body_buf = Vec::new();
let mut response_content_length: Option<usize> = None;
let mut is_chunked = false;
let mut got_first_byte = false;
let mut header_buf = Vec::with_capacity(8192);
loop {
let n = match tokio::time::timeout(READ_TIMEOUT, conn.read(&mut tmp)).await {
Ok(Ok(0)) => {
// Upstream closed — connection is no longer reusable
upstream_ok = false;
break;
}
// 15s idle timeout after first byte, 60s for initial response
let timeout = if got_first_byte {
std::time::Duration::from_secs(15)
} else {
std::time::Duration::from_secs(60)
};
let n = match tokio::time::timeout(timeout, conn.read(&mut tmp)).await {
Ok(Ok(0)) => { upstream_ok = false; break; }
Ok(Ok(n)) => n,
Ok(Err(e)) => {
debug!(domain, error = %e, "MITM: upstream read finished");
debug!(domain, error = %e, "MITM: upstream read ended");
upstream_ok = false;
break;
}
Err(_) => {
warn!(domain, "MITM: upstream read timed out after 5 minutes");
if got_first_byte {
debug!(domain, "MITM: response idle timeout (complete)");
} else {
warn!(domain, "MITM: no upstream response in 60s");
}
upstream_ok = false;
break;
}
};
got_first_byte = true;
let chunk = &tmp[..n];
// Check response headers for content-type
if !headers_parsed {
// We need to buffer until we see the end of headers
let buf = non_streaming_buf.get_or_insert_with(|| Vec::with_capacity(1024 * 64));
buf.extend_from_slice(chunk);
if let Some(_hdr_end) = find_headers_end(buf) {
// Use httparse for response header parsing
header_buf.extend_from_slice(chunk);
if let Some(_hdr_end) = find_headers_end(&header_buf) {
let mut resp_headers = [httparse::EMPTY_HEADER; 64];
let mut resp = httparse::Response::new(&mut resp_headers);
let hdr_end = match resp.parse(buf) {
let hdr_end = match resp.parse(&header_buf) {
Ok(httparse::Status::Complete(n)) => n,
_ => _hdr_end, // Fallback to manual detection
_ => _hdr_end,
};
// Detect content type and connection handling from parsed headers
let mut content_type = String::new();
for header in resp.headers.iter() {
if header.name.eq_ignore_ascii_case("content-type") {
if let Ok(val) = std::str::from_utf8(header.value) {
if val.contains("text/event-stream") {
is_streaming_response = true;
}
if let Ok(v) = std::str::from_utf8(header.value) {
content_type = v.to_string();
if v.contains("text/event-stream") { is_streaming_response = true; }
}
}
if header.name.eq_ignore_ascii_case("content-length") {
if let Ok(v) = std::str::from_utf8(header.value) {
response_content_length = v.trim().parse().ok();
}
}
if header.name.eq_ignore_ascii_case("connection") {
if let Ok(val) = std::str::from_utf8(header.value) {
if val.trim().eq_ignore_ascii_case("close") {
upstream_ok = false;
}
if let Ok(v) = std::str::from_utf8(header.value) {
if v.trim().eq_ignore_ascii_case("close") { upstream_ok = false; }
}
}
if header.name.eq_ignore_ascii_case("transfer-encoding") {
if let Ok(v) = std::str::from_utf8(header.value) {
if v.trim().eq_ignore_ascii_case("chunked") { is_chunked = true; }
}
}
}
info!(domain, streaming = is_streaming_response,
content_length = ?response_content_length,
content_type = %content_type,
status = resp.code, "MITM: got response headers");
headers_parsed = true;
if is_streaming_response {
// For streaming, parse any SSE data already in the buffer
let body_so_far = String::from_utf8_lossy(&buf[hdr_end..]);
if !body_so_far.is_empty() {
parse_streaming_chunk(&body_so_far, &mut streaming_acc);
}
// Forward the accumulated buffer to client
if let Err(e) = client.write_all(buf).await {
warn!(error = %e, "MITM: write to client failed");
break;
}
non_streaming_buf = None;
continue;
// Save body for usage parsing
response_body_buf.extend_from_slice(&header_buf[hdr_end..]);
// Forward to client immediately
if let Err(e) = client.write_all(&header_buf).await {
warn!(error = %e, "MITM: write to client failed");
break;
}
if is_streaming_response && hdr_end < header_buf.len() {
let body = String::from_utf8_lossy(&header_buf[hdr_end..]);
parse_streaming_chunk(&body, &mut streaming_acc);
}
if let Some(cl) = response_content_length {
if response_body_buf.len() >= cl { break; }
}
// Check chunked terminator in initial body
if is_chunked && has_chunked_terminator(&response_body_buf) {
debug!(domain, "MITM: chunked response complete (initial)");
break;
}
// Non-streaming: keep buffering the response body for parsing
continue;
}
continue;
}
// If streaming, parse SSE events and forward immediately
// Forward to client immediately
if let Err(e) = client.write_all(chunk).await {
warn!(error = %e, "MITM: write to client failed");
break;
}
response_body_buf.extend_from_slice(chunk);
if is_streaming_response {
let chunk_str = String::from_utf8_lossy(chunk);
parse_streaming_chunk(&chunk_str, &mut streaming_acc);
if let Err(e) = client.write_all(chunk).await {
warn!(error = %e, "MITM: write to client failed (client disconnected?)");
break;
}
} else {
// Non-streaming: keep accumulating to parse usage at the end
if let Some(ref mut buf) = non_streaming_buf {
buf.extend_from_slice(chunk);
}
let s = String::from_utf8_lossy(chunk);
parse_streaming_chunk(&s, &mut streaming_acc);
}
if let Some(cl) = response_content_length {
if response_body_buf.len() >= cl { break; }
}
if is_chunked && has_chunked_terminator(&response_body_buf) {
debug!(domain, "MITM: chunked response complete");
break;
}
}
// Forward non-streaming response all at once
if !is_streaming_response {
if let Some(ref buf) = non_streaming_buf {
if let Err(e) = client.write_all(buf).await {
warn!(error = %e, "MITM: write to client failed");
}
}
}
// Flush client
let _ = client.flush().await;
// Capture usage data
if is_streaming_response {
@@ -611,12 +705,9 @@ async fn handle_http_over_tls(
let usage = streaming_acc.into_usage();
store.record_usage(cascade_hint.as_deref(), usage).await;
}
} else if let Some(ref buf) = non_streaming_buf {
if let Some(body_start) = find_headers_end(buf) {
let body = &buf[body_start..];
if let Some(usage) = parse_non_streaming_response(body) {
store.record_usage(cascade_hint.as_deref(), usage).await;
}
} else if !response_body_buf.is_empty() {
if let Some(usage) = parse_non_streaming_response(&response_body_buf) {
store.record_usage(cascade_hint.as_deref(), usage).await;
}
}
@@ -652,6 +743,20 @@ async fn handle_passthrough(
Ok(())
}
/// Detect end of HTTP chunked transfer encoding.
/// A chunked response ends with "0\r\n\r\n" (zero-length chunk + empty trailer).
/// We check the tail of the buffer for this pattern.
fn has_chunked_terminator(body: &[u8]) -> bool {
// The minimal terminator is "0\r\n\r\n" (5 bytes)
if body.len() < 5 {
return false;
}
// Check last 7 bytes to account for possible trailing whitespace
let tail = if body.len() > 7 { &body[body.len() - 7..] } else { body };
// Look for \r\n0\r\n\r\n anywhere in the tail
tail.windows(5).any(|w| w == b"0\r\n\r\n")
}
/// Check if buffer contains a complete HTTP request (headers + full body).
/// Uses `httparse` for zero-copy, case-insensitive header parsing.
fn has_complete_http_request(buf: &[u8]) -> bool {

45
src/proto.rs
View File

@@ -62,6 +62,51 @@ pub fn varint_field(field: u32, val: u64) -> Vec<u8> {
out
}
// ─── Init metadata builder (for standalone LS stdin) ─────────────────────────
/// Build the init metadata protobuf that the LS expects on stdin at startup.
///
/// This replaces the Python snippet in `standalone-ls.sh` with proper Rust encoding.
/// Fields match what the real Antigravity extension sends to the LS.
///
/// Field layout (from binary analysis):
/// 1: api_key (string) — unique session key
/// 3: ide_name (string) — "antigravity"
/// 4: antigravity_version (string) — e.g. "1.107.0"
/// 5: ide_version (string) — e.g. "1.16.5"
/// 6: locale (string) — "en_US"
/// 10: session_id (string) — unique session identifier
/// 11: editor_name (string) — "antigravity"
/// 34: detect_and_use_proxy (varint enum) — 1 = ENABLED
pub fn build_init_metadata(
api_key: &str,
antigravity_version: &str,
ide_version: &str,
session_id: &str,
detect_and_use_proxy: u64,
) -> Vec<u8> {
let mut buf = Vec::with_capacity(128);
// Field 1: api_key
buf.extend(proto_string(1, api_key.as_bytes()));
// Field 3: ide_name
buf.extend(proto_string(3, CLIENT_NAME.as_bytes()));
// Field 4: antigravity version
buf.extend(proto_string(4, antigravity_version.as_bytes()));
// Field 5: IDE/client version
buf.extend(proto_string(5, ide_version.as_bytes()));
// Field 6: locale
buf.extend(proto_string(6, b"en_US"));
// Field 10: session_id
buf.extend(proto_string(10, session_id.as_bytes()));
// Field 11: editor_name
buf.extend(proto_string(11, CLIENT_NAME.as_bytes()));
// Field 34: detect_and_use_proxy enum (1 = ENABLED)
buf.extend(varint_field(34, detect_and_use_proxy));
buf
}
// ─── SendUserCascadeMessageRequest builder ───────────────────────────────────
/// Build the `SendUserCascadeMessageRequest` protobuf binary.

373
src/standalone.rs Normal file
View File

@@ -0,0 +1,373 @@
//! Standalone Language Server — spawn and lifecycle management.
//!
//! Launches an isolated LS instance as a child process that the proxy fully owns.
//! The standalone LS shares auth via the main extension server but has its own
//! HTTPS port, data directory, and cascade space. This means the real LS (the
//! one powering the user's coding session) is never touched.
use crate::constants;
use crate::proto;
use std::io::Write;
use std::net::TcpListener;
use std::process::{Child, Command, Stdio};
use tokio::time::{sleep, Duration};
use tracing::{debug, info};
/// Default path to the LS binary.
const LS_BINARY_PATH: &str =
"/usr/share/antigravity/resources/app/extensions/antigravity/bin/language_server_linux_x64";
/// App root for ANTIGRAVITY_EDITOR_APP_ROOT env var.
const APP_ROOT: &str = "/usr/share/antigravity/resources/app";
/// Data directory for the standalone LS.
const DATA_DIR: &str = "/tmp/antigravity-standalone";
/// System user for UID-scoped iptables isolation.
const LS_USER: &str = "antigravity-ls";
/// A running standalone LS process.
pub struct StandaloneLS {
child: Child,
pub port: u16,
pub csrf: String,
}
/// Config needed from the real (main) LS to bootstrap the standalone one.
pub struct MainLSConfig {
pub extension_server_port: String,
pub csrf: String,
}
/// Optional MITM proxy config for the standalone LS.
pub struct StandaloneMitmConfig {
pub proxy_addr: String, // e.g. "http://127.0.0.1:8742"
pub ca_cert_path: String, // path to MITM CA .pem
}
impl StandaloneLS {
/// Spawn a standalone LS process.
///
/// Discovers the main LS's extension server port and CSRF token,
/// picks a free port, builds init metadata, and launches the binary.
///
/// If `mitm_config` is provided, sets HTTPS_PROXY and SSL_CERT_FILE
/// so the LS routes LLM API calls through the MITM proxy.
pub fn spawn(
main_config: &MainLSConfig,
mitm_config: Option<&StandaloneMitmConfig>,
) -> Result<Self, String> {
let port = find_free_port()?;
let ts = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
// Build init metadata protobuf
let api_key = format!("standalone-api-key-{ts}");
let session_id = format!("standalone-session-{ts}");
let metadata = proto::build_init_metadata(
&api_key,
constants::antigravity_version(),
constants::client_version(),
&session_id,
1, // DETECT_AND_USE_PROXY_ENABLED
);
// Setup data dir (mode 1777 so both current user and antigravity-ls can write)
let gemini_dir = format!("{DATA_DIR}/.gemini");
std::fs::create_dir_all(&gemini_dir)
.map_err(|e| format!("Failed to create standalone data dir: {e}"))?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let _ = std::fs::set_permissions(DATA_DIR, std::fs::Permissions::from_mode(0o1777));
let _ = std::fs::set_permissions(&gemini_dir, std::fs::Permissions::from_mode(0o1777));
}
// LS args — mirrors standalone-ls.sh but with correct params
let args = vec![
"-enable_lsp".to_string(),
"-extension_server_port".to_string(),
main_config.extension_server_port.clone(),
"-csrf_token".to_string(),
main_config.csrf.clone(),
"-server_port".to_string(),
port.to_string(),
"-workspace_id".to_string(),
format!("standalone_{ts}"),
"-cloud_code_endpoint".to_string(),
"https://daily-cloudcode-pa.googleapis.com".to_string(),
"-app_data_dir".to_string(),
"antigravity-standalone".to_string(),
"-gemini_dir".to_string(),
gemini_dir,
];
info!(port, "Spawning standalone LS");
debug!(?args, "LS args");
// Build env vars for the LS process
let mut env_vars: Vec<(String, String)> = vec![
("ANTIGRAVITY_EDITOR_APP_ROOT".into(), APP_ROOT.into()),
];
// If MITM is enabled, add SSL + proxy env vars
if let Some(mitm) = mitm_config {
// Go's SSL_CERT_FILE replaces the entire system cert pool, so we
// need a combined bundle: system CAs + our MITM CA
// Write to /tmp — accessible by antigravity-ls user
// (user's ~/.config/ is not traversable by other UIDs)
let combined_ca_path = "/tmp/antigravity-mitm-combined-ca.pem".to_string();
let system_ca = std::fs::read_to_string("/etc/ssl/certs/ca-certificates.crt")
.unwrap_or_default();
let mitm_ca = std::fs::read_to_string(&mitm.ca_cert_path)
.map_err(|e| format!("Failed to read MITM CA cert: {e}"))?;
std::fs::write(&combined_ca_path, format!("{system_ca}\n{mitm_ca}"))
.map_err(|e| format!("Failed to write combined CA bundle: {e}"))?;
// Make readable by antigravity-ls user
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let _ = std::fs::set_permissions(
&combined_ca_path,
std::fs::Permissions::from_mode(0o644),
);
}
info!(
proxy = %mitm.proxy_addr,
ca = %combined_ca_path,
"Setting MITM env vars on standalone LS (combined CA bundle)"
);
env_vars.push(("SSL_CERT_FILE".into(), combined_ca_path));
env_vars.push(("SSL_CERT_DIR".into(), "/dev/null".into()));
env_vars.push(("NODE_EXTRA_CA_CERTS".into(), mitm.ca_cert_path.clone()));
}
// Check if 'antigravity-ls' user exists for UID-scoped iptables isolation
let use_sudo = has_ls_user();
let mut cmd = if use_sudo {
info!("Using UID isolation: spawning LS as 'antigravity-ls' user");
// Build: sudo -n -u antigravity-ls -- /usr/bin/env VAR=val ... LS_BINARY args...
let mut c = Command::new("sudo");
c.args(["-n", "-u", LS_USER, "--", "/usr/bin/env"]);
// Pass env vars as key=value args to /usr/bin/env
for (k, v) in &env_vars {
c.arg(format!("{k}={v}"));
}
c.arg(LS_BINARY_PATH);
c.args(&args);
c
} else {
debug!("No 'antigravity-ls' user found, spawning LS as current user");
let mut c = Command::new(LS_BINARY_PATH);
c.args(&args);
for (k, v) in &env_vars {
c.env(k, v);
}
c
};
cmd.stdin(Stdio::piped())
.stdout(Stdio::null())
.stderr(Stdio::null());
let mut child = cmd
.spawn()
.map_err(|e| format!("Failed to spawn LS binary: {e}"))?;
// Feed init metadata via stdin, then close it
if let Some(mut stdin) = child.stdin.take() {
stdin
.write_all(&metadata)
.map_err(|e| format!("Failed to write init metadata to stdin: {e}"))?;
// stdin drops here → EOF
}
info!(pid = child.id(), port, "Standalone LS spawned");
Ok(StandaloneLS {
child,
port,
csrf: main_config.csrf.clone(),
})
}
/// Wait for the standalone LS to be ready (accepting TCP connections).
///
/// Retries up to `max_attempts` times with a 1-second delay between each.
pub async fn wait_ready(&self, max_attempts: u32) -> Result<(), String> {
info!(port = self.port, "Waiting for standalone LS to be ready...");
for attempt in 1..=max_attempts {
sleep(Duration::from_secs(1)).await;
// Simple TCP connect check — if the LS is listening, it's ready
match tokio::net::TcpStream::connect(format!("127.0.0.1:{}", self.port)).await {
Ok(_) => {
info!(attempt, "Standalone LS is ready (accepting connections)");
return Ok(());
}
Err(e) => {
debug!(attempt, error = %e, "LS not ready yet");
}
}
}
Err(format!(
"Standalone LS failed to become ready after {max_attempts} attempts on port {}",
self.port
))
}
/// Check if the child process is still running.
#[allow(dead_code)]
pub fn is_alive(&mut self) -> bool {
matches!(self.child.try_wait(), Ok(None))
}
/// Kill the standalone LS process.
pub fn kill(&mut self) {
info!("Killing standalone LS");
let _ = self.child.kill();
let _ = self.child.wait();
}
}
impl Drop for StandaloneLS {
fn drop(&mut self) {
self.kill();
}
}
/// Discover only the extension_server_port and csrf_token from the running main LS.
///
/// This does NOT discover the HTTPS port — we don't need to talk to the real LS,
/// only steal its extension server connection info.
pub fn discover_main_ls_config() -> Result<MainLSConfig, String> {
let pid = find_main_ls_pid()?;
let cmdline = std::fs::read(format!("/proc/{pid}/cmdline"))
.map_err(|e| format!("Can't read cmdline for PID {pid}: {e}"))?;
let args: Vec<&[u8]> = cmdline.split(|&b| b == 0).collect();
let mut csrf = String::new();
let mut ext_port = String::new();
for (i, arg) in args.iter().enumerate() {
if let Ok(s) = std::str::from_utf8(arg) {
match s {
"--csrf_token" | "-csrf_token" => {
if let Some(next) = args.get(i + 1) {
if let Ok(val) = std::str::from_utf8(next) {
csrf = val.to_string();
}
}
}
"--extension_server_port" | "-extension_server_port" => {
if let Some(next) = args.get(i + 1) {
if let Ok(val) = std::str::from_utf8(next) {
ext_port = val.to_string();
}
}
}
_ => {}
}
}
}
if csrf.is_empty() {
return Err("Could not find CSRF token from main LS".to_string());
}
if ext_port.is_empty() {
return Err("Could not find extension_server_port from main LS".to_string());
}
info!(
pid,
ext_port,
csrf_len = csrf.len(),
"Discovered main LS config"
);
Ok(MainLSConfig {
extension_server_port: ext_port,
csrf,
})
}
/// Find the PID of the main (real) LS process.
///
/// Checks `/proc/<pid>/exe` to ensure we find the actual LS binary,
/// not bash scripts that happen to mention `language_server_linux` in their args.
fn find_main_ls_pid() -> Result<String, String> {
let proc = std::path::Path::new("/proc");
if !proc.exists() {
return Err("No /proc filesystem".to_string());
}
let entries = std::fs::read_dir(proc)
.map_err(|e| format!("Cannot read /proc: {e}"))?;
for entry in entries.flatten() {
let name = entry.file_name();
let name_str = name.to_string_lossy();
// Only numeric dirs (PIDs)
if !name_str.chars().all(|c| c.is_ascii_digit()) {
continue;
}
let exe_link = entry.path().join("exe");
if let Ok(target) = std::fs::read_link(&exe_link) {
let target_str = target.to_string_lossy().to_string();
let target_clean = target_str.trim_end_matches(" (deleted)");
// Must be the actual LS binary, not a bash script
if target_clean.contains("language_server_linux")
|| target_clean.contains("antigravity-language-server")
{
return Ok(name_str.to_string());
}
}
}
Err("No main LS process found — Antigravity must be running".to_string())
}
/// Find a free TCP port by binding to port 0.
fn find_free_port() -> Result<u16, String> {
let listener =
TcpListener::bind("127.0.0.1:0").map_err(|e| format!("Failed to bind for port: {e}"))?;
listener
.local_addr()
.map(|a| a.port())
.map_err(|e| format!("Failed to get port: {e}"))
}
/// Check if the dedicated LS system user exists.
///
/// When the user exists, the proxy spawns the LS as that UID so iptables
/// can scope the :443 redirect to only the standalone LS process.
fn has_ls_user() -> bool {
Command::new("id")
.args(["-u", LS_USER])
.stdout(Stdio::null())
.stderr(Stdio::null())
.status()
.map(|s| s.success())
.unwrap_or(false)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_find_free_port() {
let port = find_free_port().unwrap();
assert!(port > 0);
// Port should be available — try binding to it
let listener = TcpListener::bind(format!("127.0.0.1:{port}"));
assert!(listener.is_ok(), "Port {port} should be free");
}
}