feat: MITM interception for standalone LS with UID isolation
- 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:
99
src/main.rs
99
src/main.rs
@@ -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");
|
||||
|
||||
Reference in New Issue
Block a user