447 lines
16 KiB
Rust
447 lines
16 KiB
Rust
//! Antigravity OpenAI Proxy — Rust edition v3 (stealth hardened).
|
|
//!
|
|
//! Single-binary replacement for server.py. BoringSSL TLS impersonation,
|
|
//! byte-exact protobuf encoding, Chrome header fingerprinting, cascade
|
|
//! session management, warmup + heartbeat lifecycle mimicry.
|
|
|
|
mod api;
|
|
mod backend;
|
|
mod constants;
|
|
mod mitm;
|
|
mod proto;
|
|
mod quota;
|
|
mod session;
|
|
mod standalone;
|
|
mod warmup;
|
|
|
|
use api::AppState;
|
|
use backend::Backend;
|
|
use clap::Parser;
|
|
use session::SessionManager;
|
|
use std::sync::Arc;
|
|
use tracing::{info, warn};
|
|
|
|
use mitm::store::MitmStore;
|
|
|
|
#[derive(Parser)]
|
|
#[command(name = "antigravity-proxy", about = "Antigravity OpenAI Proxy (stealth)")]
|
|
struct Cli {
|
|
/// Port to listen on
|
|
#[arg(long, default_value_t = 8741)]
|
|
port: u16,
|
|
|
|
/// Enable info-level logging (-v)
|
|
#[arg(short, long)]
|
|
verbose: bool,
|
|
|
|
/// Enable debug-level logging (-d)
|
|
#[arg(short, long)]
|
|
debug: bool,
|
|
|
|
/// Disable the MITM proxy (no API interception)
|
|
#[arg(long)]
|
|
no_mitm: bool,
|
|
|
|
/// MITM proxy port (default: 8742, matches wrapper script)
|
|
#[arg(long, default_value_t = 8742)]
|
|
mitm_port: u16,
|
|
|
|
/// Disable standalone LS — attach to the real running LS instead
|
|
#[arg(long)]
|
|
no_standalone: bool,
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() {
|
|
// Ignore SIGPIPE — prevents instant death when piped through tee/grep
|
|
#[cfg(unix)]
|
|
{
|
|
use tokio::signal::unix::{signal, SignalKind};
|
|
let mut sigpipe = signal(SignalKind::pipe()).expect("failed to install SIGPIPE handler");
|
|
tokio::spawn(async move {
|
|
loop {
|
|
sigpipe.recv().await;
|
|
// Silently ignore SIGPIPE
|
|
}
|
|
});
|
|
}
|
|
|
|
// Install rustls CryptoProvider early — prevents panic under concurrent load
|
|
let _ = rustls::crypto::ring::default_provider().install_default();
|
|
|
|
let cli = Cli::parse();
|
|
|
|
// Flag > env var > default (warn)
|
|
let log_level = if cli.debug {
|
|
"debug"
|
|
} else if cli.verbose {
|
|
"info"
|
|
} else {
|
|
// Fall back to RUST_LOG env, or warn-only
|
|
""
|
|
};
|
|
|
|
let filter = if log_level.is_empty() {
|
|
tracing_subscriber::EnvFilter::try_from_default_env()
|
|
.unwrap_or_else(|_| "warn".into())
|
|
} else {
|
|
tracing_subscriber::EnvFilter::new(log_level)
|
|
};
|
|
|
|
tracing_subscriber::fmt()
|
|
.with_env_filter(filter)
|
|
.init();
|
|
|
|
// ── Step 1: Bind main port FIRST (fail fast, before spawning anything) ────
|
|
let addr = format!("127.0.0.1:{}", cli.port);
|
|
let listener = match tokio::net::TcpListener::bind(&addr).await {
|
|
Ok(l) => l,
|
|
Err(e) => {
|
|
eprintln!("Fatal: cannot bind to {addr}: {e}");
|
|
eprintln!("Hint: kill $(lsof -ti:{}) 2>/dev/null", cli.port);
|
|
std::process::exit(1);
|
|
}
|
|
};
|
|
|
|
// ── Step 2: Backend discovery (or standalone LS spawn) ─────────────────────
|
|
let standalone_ls = if !cli.no_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);
|
|
}
|
|
}
|
|
});
|
|
|
|
let (pid, https_port, csrf, token) = backend.info().await;
|
|
|
|
// ── Step 3: MITM proxy (after port is secured) ────────────────────────────
|
|
let mitm_store = MitmStore::new();
|
|
let (mitm_port_actual, mitm_handle) = if !cli.no_mitm {
|
|
let data_dir = dirs_data_dir();
|
|
match mitm::ca::MitmCa::load_or_generate(&data_dir) {
|
|
Ok(ca) => {
|
|
let ca = Arc::new(ca);
|
|
let ca_pem = ca.ca_pem_path.display().to_string();
|
|
let config = mitm::proxy::MitmConfig {
|
|
port: cli.mitm_port,
|
|
modify_requests: true,
|
|
};
|
|
match mitm::proxy::run(ca, mitm_store.clone(), config).await {
|
|
Ok((port, handle)) => {
|
|
info!(port, ca = %ca_pem, "MITM proxy started");
|
|
// Write actual port to file for wrapper script discovery
|
|
let port_file = data_dir.join("mitm-port");
|
|
if let Err(e) = std::fs::write(&port_file, port.to_string()) {
|
|
warn!("Failed to write MITM port file: {e}");
|
|
}
|
|
(Some((port, ca_pem)), Some(handle))
|
|
}
|
|
Err(e) => {
|
|
warn!("MITM proxy failed to start: {e}");
|
|
(None, None)
|
|
}
|
|
}
|
|
}
|
|
Err(e) => {
|
|
warn!("MITM CA generation failed: {e}");
|
|
(None, None)
|
|
}
|
|
}
|
|
} else {
|
|
info!("MITM proxy disabled (--no-mitm)");
|
|
(None, None)
|
|
};
|
|
|
|
// ── Step 4: Warmup + heartbeat ────────────────────────────────────────────
|
|
warmup::warmup_sequence(&backend).await;
|
|
let heartbeat_handle = warmup::start_heartbeat(Arc::clone(&backend));
|
|
|
|
// ── Step 4b: Quota monitor ────────────────────────────────────────────────
|
|
let quota_store = quota::QuotaStore::new();
|
|
quota_store.clone().start_polling(Arc::clone(&backend));
|
|
info!("Quota monitor started (polling every 60s)");
|
|
|
|
let state = Arc::new(AppState {
|
|
backend,
|
|
sessions: SessionManager::new(),
|
|
mitm_store,
|
|
quota_store,
|
|
});
|
|
|
|
// 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.no_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 {
|
|
warn!("Periodic refresh failed: {e}");
|
|
}
|
|
}
|
|
});
|
|
|
|
// ── Step 5: Start serving ─────────────────────────────────────────────────
|
|
let app = api::router(state.clone());
|
|
|
|
print_banner(cli.port, &pid, &https_port, &csrf, &token, &mitm_port_actual, is_standalone);
|
|
info!("Listening on http://{addr}");
|
|
|
|
axum::serve(listener, app)
|
|
.with_graceful_shutdown(shutdown_signal())
|
|
.await
|
|
.expect("server error");
|
|
|
|
// ── Cleanup: abort all background tasks ───────────────────────────────────
|
|
heartbeat_handle.abort();
|
|
refresh_handle.abort();
|
|
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"));
|
|
eprintln!(" \x1b[1;32m✓ Server shutdown complete\x1b[0m\n");
|
|
info!("Server shutdown complete");
|
|
}
|
|
|
|
/// Wait for SIGINT (Ctrl+C) or SIGTERM for graceful shutdown.
|
|
async fn shutdown_signal() {
|
|
let ctrl_c = async {
|
|
tokio::signal::ctrl_c()
|
|
.await
|
|
.expect("failed to install Ctrl+C handler");
|
|
};
|
|
|
|
#[cfg(unix)]
|
|
let terminate = async {
|
|
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
|
|
.expect("failed to install SIGTERM handler")
|
|
.recv()
|
|
.await;
|
|
};
|
|
|
|
#[cfg(not(unix))]
|
|
let terminate = std::future::pending::<()>();
|
|
|
|
tokio::select! {
|
|
_ = ctrl_c => {
|
|
eprintln!("\n \x1b[1;33m⚡ Shutting down gracefully...\x1b[0m");
|
|
info!("Received SIGINT, shutting down...");
|
|
},
|
|
_ = terminate => {
|
|
eprintln!("\n \x1b[1;33m⚡ Received SIGTERM, shutting down...\x1b[0m");
|
|
info!("Received SIGTERM, shutting down...");
|
|
},
|
|
}
|
|
}
|
|
|
|
fn print_banner(port: u16, pid: &str, https_port: &str, csrf: &str, token: &str, mitm: &Option<(u16, String)>, is_standalone: bool) {
|
|
let chrome_major = &*constants::CHROME_MAJOR;
|
|
let ver = crate::constants::antigravity_version();
|
|
|
|
println!();
|
|
println!(" \x1b[1;35m>> antigravity-proxy\x1b[0m \x1b[2mv{ver}\x1b[0m");
|
|
println!(" \x1b[2m────────────────────────────────────────────────\x1b[0m");
|
|
println!();
|
|
println!(" \x1b[1mcore\x1b[0m");
|
|
println!(" \x1b[36m tls\x1b[0m BoringSSL (Chrome {chrome_major})");
|
|
println!(" \x1b[36m listen\x1b[0m http://127.0.0.1:{port}");
|
|
println!(" \x1b[36m ls pid\x1b[0m {pid}");
|
|
println!(" \x1b[36m https\x1b[0m :{https_port}");
|
|
println!(" \x1b[36m csrf\x1b[0m {csrf}");
|
|
println!(" \x1b[36m oauth\x1b[0m {token}");
|
|
println!();
|
|
|
|
// MITM section
|
|
if let Some((mitm_port, ca_path)) = mitm {
|
|
println!(" \x1b[1mmitm\x1b[0m");
|
|
println!(" \x1b[36m proxy\x1b[0m 127.0.0.1:{mitm_port}");
|
|
println!(" \x1b[36m ca cert\x1b[0m {ca_path}");
|
|
|
|
// Check if wrapper is installed
|
|
if is_standalone {
|
|
println!(" \x1b[36m wrapper\x1b[0m \x1b[32miptables (standalone)\x1b[0m");
|
|
} else if check_wrapper_installed() {
|
|
println!(" \x1b[36m wrapper\x1b[0m \x1b[32minstalled\x1b[0m");
|
|
} else {
|
|
println!(" \x1b[36m wrapper\x1b[0m \x1b[33mnot installed\x1b[0m");
|
|
}
|
|
println!();
|
|
} else {
|
|
println!(" \x1b[1mmitm\x1b[0m \x1b[33mdisabled\x1b[0m");
|
|
println!();
|
|
}
|
|
|
|
// Routes
|
|
println!(" \x1b[1mroutes\x1b[0m");
|
|
println!(" \x1b[33m POST\x1b[0m /v1/responses");
|
|
println!(" \x1b[33m POST\x1b[0m /v1/chat/completions");
|
|
println!(" \x1b[33m POST\x1b[0m /v1/gemini");
|
|
println!(" \x1b[32m GET \x1b[0m /v1/models");
|
|
println!(" \x1b[32m GET \x1b[0m /v1/sessions");
|
|
println!(" \x1b[31m DEL \x1b[0m /v1/sessions/:id");
|
|
println!(" \x1b[33m POST\x1b[0m /v1/token");
|
|
println!(" \x1b[32m GET \x1b[0m /v1/usage");
|
|
println!(" \x1b[32m GET \x1b[0m /v1/quota");
|
|
println!(" \x1b[32m GET \x1b[0m /health");
|
|
println!();
|
|
|
|
// Status line
|
|
let mitm_tag = if mitm.is_some() { "\x1b[32mmitm\x1b[0m" } else { "\x1b[31mmitm\x1b[0m" };
|
|
println!(" \x1b[2mstealth:\x1b[0m \x1b[32mwarmup\x1b[0m \x1b[32mheartbeat\x1b[0m \x1b[32mjitter\x1b[0m {mitm_tag}");
|
|
println!();
|
|
|
|
// Setup hints
|
|
if let Some((mitm_port, ca_path)) = mitm {
|
|
if is_standalone {
|
|
// Standalone mode uses iptables UID isolation — no wrapper needed
|
|
} else if !check_wrapper_installed() {
|
|
println!(" \x1b[1;33m[!]\x1b[0m mitm wrapper not installed");
|
|
println!(" \x1b[2mrun:\x1b[0m ./scripts/mitm-wrapper.sh install");
|
|
println!(" \x1b[2mor:\x1b[0m HTTPS_PROXY=http://127.0.0.1:{mitm_port}");
|
|
println!(" NODE_EXTRA_CA_CERTS={ca_path}");
|
|
println!();
|
|
}
|
|
}
|
|
|
|
if token == "NOT SET" {
|
|
println!(" \x1b[1;33m[!]\x1b[0m no oauth token");
|
|
println!(" export ANTIGRAVITY_OAUTH_TOKEN=ya29.xxx");
|
|
println!(" curl -X POST http://127.0.0.1:{port}/v1/token -d '{{\"token\":\"ya29.xxx\"}}'");
|
|
println!(" echo 'ya29.xxx' > ~/.config/antigravity-proxy-token");
|
|
println!();
|
|
}
|
|
}
|
|
|
|
/// Check if the MITM wrapper is installed by looking for the .real backup file
|
|
/// next to the LS binary. Uses /proc to find the real LS path dynamically.
|
|
fn check_wrapper_installed() -> bool {
|
|
// Find the LS binary path from known PID or by scanning /proc
|
|
if let Some(ls_path) = find_ls_binary_path() {
|
|
let real_path = format!("{ls_path}.real");
|
|
return std::path::Path::new(&real_path).exists();
|
|
}
|
|
false
|
|
}
|
|
|
|
/// Find the LS binary path by reading /proc/<pid>/exe for known language server processes.
|
|
fn find_ls_binary_path() -> Option<String> {
|
|
// Try all running processes, look for ones that look like the LS
|
|
let proc = std::path::Path::new("/proc");
|
|
if !proc.exists() {
|
|
return None;
|
|
}
|
|
|
|
if let Ok(entries) = std::fs::read_dir(proc) {
|
|
for entry in entries.flatten() {
|
|
let name = entry.file_name();
|
|
let name_str = name.to_string_lossy();
|
|
// Only look at 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();
|
|
// Strip " (deleted)" suffix from unlinked binaries
|
|
let target_clean = target_str.trim_end_matches(" (deleted)");
|
|
// Match any binary that looks like the Antigravity LS
|
|
if target_clean.contains("language_server_linux")
|
|
|| target_clean.contains("antigravity-language-server")
|
|
{
|
|
// Strip .real suffix — if the wrapper exec'd the backup, we want the base name
|
|
let path = target_clean.trim_end_matches(".real");
|
|
return Some(path.to_string());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
None
|
|
}
|
|
|
|
/// Get the data directory for storing MITM CA cert/key.
|
|
fn dirs_data_dir() -> std::path::PathBuf {
|
|
let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
|
|
std::path::PathBuf::from(home).join(".config").join("antigravity-proxy")
|
|
}
|