mirror of
https://github.com/therealaleph/MasterHttpRelayVPN-RUST.git
synced 2026-05-18 07:34:36 +03:00
feat: tunnel-node service + CodeFull.gs (#93)
Standalone Rust/axum HTTP server + Apps Script-side CodeFull.gs for users who want to deploy a remote tunnel node. All new files; no changes to the main Rust crate. This is part 1 of 3 of the full-tunnel feature — it adds scaffolding that users can opt into once the Rust-side Mode::Full lands in #94.
This commit is contained in:
vahidlazio
@@ -0,0 +1 @@
|
||||
/target
|
||||
Generated
+1043
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,27 @@
|
||||
[package]
|
||||
name = "mhrv-tunnel-node"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
description = "HTTP tunnel bridge for MasterHttpRelayVPN full mode — bridges HTTP tunnel requests to real TCP connections"
|
||||
|
||||
[[bin]]
|
||||
name = "tunnel-node"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
axum = "0.7"
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "macros", "net", "time", "io-util", "signal", "sync"] }
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
base64 = "0.22"
|
||||
uuid = { version = "1", features = ["v4"] }
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
||||
flate2 = "1"
|
||||
|
||||
[profile.release]
|
||||
panic = "abort"
|
||||
codegen-units = 1
|
||||
lto = true
|
||||
opt-level = 3
|
||||
strip = true
|
||||
@@ -0,0 +1,12 @@
|
||||
FROM rust:1.85-slim AS builder
|
||||
WORKDIR /app
|
||||
COPY Cargo.toml ./
|
||||
COPY src/ ./src/
|
||||
RUN cargo build --release --bin tunnel-node
|
||||
|
||||
FROM debian:bookworm-slim
|
||||
RUN apt-get update && apt-get install -y ca-certificates && rm -rf /var/lib/apt/lists/*
|
||||
COPY --from=builder /app/target/release/tunnel-node /usr/local/bin/
|
||||
ENV PORT=8080
|
||||
EXPOSE 8080
|
||||
CMD ["tunnel-node"]
|
||||
@@ -0,0 +1,81 @@
|
||||
# Tunnel Node
|
||||
|
||||
HTTP tunnel bridge server for MasterHttpRelayVPN "full" mode. Bridges HTTP tunnel requests (from Apps Script) to real TCP connections.
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
Phone → mhrv-rs → [domain-fronted TLS] → Apps Script → [HTTP] → Tunnel Node → [real TCP] → Internet
|
||||
```
|
||||
|
||||
The tunnel node manages persistent TCP sessions. Each session is a real TCP connection to a destination server. Data flows through a JSON protocol:
|
||||
|
||||
- **connect** — open TCP to host:port, return session ID
|
||||
- **data** — write client data, return server response
|
||||
- **close** — tear down session
|
||||
- **batch** — process multiple ops in one HTTP request (reduces round trips)
|
||||
|
||||
## Deployment
|
||||
|
||||
### Cloud Run
|
||||
|
||||
```bash
|
||||
cd tunnel-node
|
||||
gcloud run deploy tunnel-node \
|
||||
--source . \
|
||||
--region us-central1 \
|
||||
--allow-unauthenticated \
|
||||
--set-env-vars TUNNEL_AUTH_KEY=$(openssl rand -hex 24) \
|
||||
--memory 256Mi \
|
||||
--cpu 1 \
|
||||
--max-instances 1
|
||||
```
|
||||
|
||||
### Docker (any VPS)
|
||||
|
||||
```bash
|
||||
cd tunnel-node
|
||||
docker build -t tunnel-node .
|
||||
docker run -p 8080:8080 -e TUNNEL_AUTH_KEY=your-secret tunnel-node
|
||||
```
|
||||
|
||||
### Direct binary
|
||||
|
||||
```bash
|
||||
cd tunnel-node
|
||||
cargo build --release
|
||||
TUNNEL_AUTH_KEY=your-secret PORT=8080 ./target/release/tunnel-node
|
||||
```
|
||||
|
||||
## Environment Variables
|
||||
|
||||
| Variable | Required | Default | Description |
|
||||
|----------|----------|---------|-------------|
|
||||
| `TUNNEL_AUTH_KEY` | Yes | `changeme` | Shared secret — must match `TUNNEL_AUTH_KEY` in CodeFull.gs |
|
||||
| `PORT` | No | `8080` | Listen port (Cloud Run sets this automatically) |
|
||||
|
||||
## Protocol
|
||||
|
||||
### Single op: `POST /tunnel`
|
||||
|
||||
```json
|
||||
{"k":"auth","op":"connect","host":"example.com","port":443}
|
||||
{"k":"auth","op":"data","sid":"uuid","data":"base64"}
|
||||
{"k":"auth","op":"close","sid":"uuid"}
|
||||
```
|
||||
|
||||
### Batch: `POST /tunnel/batch`
|
||||
|
||||
```json
|
||||
{
|
||||
"k": "auth",
|
||||
"ops": [
|
||||
{"op":"data","sid":"uuid1","d":"base64"},
|
||||
{"op":"data","sid":"uuid2","d":"base64"},
|
||||
{"op":"close","sid":"uuid3"}
|
||||
]
|
||||
}
|
||||
→ {"r": [{...}, {...}, {...}]}
|
||||
```
|
||||
|
||||
### Health check: `GET /health` → `ok`
|
||||
@@ -0,0 +1,529 @@
|
||||
//! HTTP Tunnel Node for MasterHttpRelayVPN "full" mode.
|
||||
//!
|
||||
//! Bridges HTTP tunnel requests (from Apps Script) to real TCP connections.
|
||||
//! Supports both single-op (`POST /tunnel`) and batch (`POST /tunnel/batch`)
|
||||
//! modes. Batch mode processes all active sessions in one HTTP round trip,
|
||||
//! dramatically reducing the number of Apps Script calls.
|
||||
//!
|
||||
//! Env vars:
|
||||
//! TUNNEL_AUTH_KEY — shared secret (required)
|
||||
//! PORT — listen port (default 8080, Cloud Run sets this)
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use axum::body::Bytes;
|
||||
use axum::extract::State;
|
||||
use axum::http::{header, StatusCode};
|
||||
use axum::response::IntoResponse;
|
||||
use axum::{routing::post, Json, Router};
|
||||
use base64::engine::general_purpose::STANDARD as B64;
|
||||
use base64::Engine;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Session
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
struct SessionInner {
|
||||
writer: Mutex<OwnedWriteHalf>,
|
||||
read_buf: Mutex<Vec<u8>>,
|
||||
eof: AtomicBool,
|
||||
last_active: Mutex<Instant>,
|
||||
}
|
||||
|
||||
struct ManagedSession {
|
||||
inner: Arc<SessionInner>,
|
||||
reader_handle: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
async fn create_session(host: &str, port: u16) -> std::io::Result<ManagedSession> {
|
||||
let addr = format!("{}:{}", host, port);
|
||||
let stream = tokio::time::timeout(Duration::from_secs(10), TcpStream::connect(&addr))
|
||||
.await
|
||||
.map_err(|_| std::io::Error::new(std::io::ErrorKind::TimedOut, "connect timeout"))??;
|
||||
let _ = stream.set_nodelay(true);
|
||||
let (reader, writer) = stream.into_split();
|
||||
|
||||
let inner = Arc::new(SessionInner {
|
||||
writer: Mutex::new(writer),
|
||||
read_buf: Mutex::new(Vec::with_capacity(32768)),
|
||||
eof: AtomicBool::new(false),
|
||||
last_active: Mutex::new(Instant::now()),
|
||||
});
|
||||
|
||||
let inner_ref = inner.clone();
|
||||
let reader_handle = tokio::spawn(reader_task(reader, inner_ref));
|
||||
|
||||
Ok(ManagedSession { inner, reader_handle })
|
||||
}
|
||||
|
||||
async fn reader_task(mut reader: OwnedReadHalf, session: Arc<SessionInner>) {
|
||||
let mut buf = vec![0u8; 65536];
|
||||
loop {
|
||||
match reader.read(&mut buf).await {
|
||||
Ok(0) => { session.eof.store(true, Ordering::Release); break; }
|
||||
Ok(n) => { session.read_buf.lock().await.extend_from_slice(&buf[..n]); }
|
||||
Err(_) => { session.eof.store(true, Ordering::Release); break; }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Drain whatever is currently buffered — no waiting.
|
||||
/// Used by batch mode where we poll frequently.
|
||||
async fn drain_now(session: &SessionInner) -> (Vec<u8>, bool) {
|
||||
let mut buf = session.read_buf.lock().await;
|
||||
let data = std::mem::take(&mut *buf);
|
||||
let eof = session.eof.load(Ordering::Acquire);
|
||||
(data, eof)
|
||||
}
|
||||
|
||||
/// Wait for response data with drain window. Used by single-op mode.
|
||||
async fn wait_and_drain(session: &SessionInner, max_wait: Duration) -> (Vec<u8>, bool) {
|
||||
let deadline = Instant::now() + max_wait;
|
||||
let mut prev_len = 0usize;
|
||||
let mut last_growth = Instant::now();
|
||||
let mut ever_had_data = false;
|
||||
|
||||
loop {
|
||||
let (cur_len, is_eof) = {
|
||||
let buf = session.read_buf.lock().await;
|
||||
(buf.len(), session.eof.load(Ordering::Acquire))
|
||||
};
|
||||
if cur_len > prev_len {
|
||||
last_growth = Instant::now();
|
||||
prev_len = cur_len;
|
||||
ever_had_data = true;
|
||||
}
|
||||
if is_eof { break; }
|
||||
if Instant::now() >= deadline { break; }
|
||||
if ever_had_data && last_growth.elapsed() > Duration::from_millis(100) { break; }
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
|
||||
let mut buf = session.read_buf.lock().await;
|
||||
let data = std::mem::take(&mut *buf);
|
||||
let eof = session.eof.load(Ordering::Acquire);
|
||||
(data, eof)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// App state
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[derive(Clone)]
|
||||
struct AppState {
|
||||
sessions: Arc<Mutex<HashMap<String, ManagedSession>>>,
|
||||
auth_key: String,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Protocol types — single op (backward compat)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct TunnelRequest {
|
||||
k: String,
|
||||
op: String,
|
||||
#[serde(default)] host: Option<String>,
|
||||
#[serde(default)] port: Option<u16>,
|
||||
#[serde(default)] sid: Option<String>,
|
||||
#[serde(default)] data: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Clone)]
|
||||
struct TunnelResponse {
|
||||
#[serde(skip_serializing_if = "Option::is_none")] sid: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")] d: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")] eof: Option<bool>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")] e: Option<String>,
|
||||
}
|
||||
|
||||
impl TunnelResponse {
|
||||
fn error(msg: impl Into<String>) -> Self {
|
||||
Self { sid: None, d: None, eof: None, e: Some(msg.into()) }
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Protocol types — batch
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct BatchRequest {
|
||||
k: String,
|
||||
ops: Vec<BatchOp>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct BatchOp {
|
||||
op: String,
|
||||
#[serde(default)] sid: Option<String>,
|
||||
#[serde(default)] host: Option<String>,
|
||||
#[serde(default)] port: Option<u16>,
|
||||
#[serde(default)] d: Option<String>, // base64 data
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct BatchResponse {
|
||||
r: Vec<TunnelResponse>,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Single-op handler (backward compat)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async fn handle_tunnel(
|
||||
State(state): State<AppState>,
|
||||
Json(req): Json<TunnelRequest>,
|
||||
) -> Json<TunnelResponse> {
|
||||
if req.k != state.auth_key {
|
||||
return Json(TunnelResponse::error("unauthorized"));
|
||||
}
|
||||
match req.op.as_str() {
|
||||
"connect" => Json(handle_connect(&state, req.host, req.port).await),
|
||||
"data" => Json(handle_data_single(&state, req.sid, req.data).await),
|
||||
"close" => Json(handle_close(&state, req.sid).await),
|
||||
other => Json(TunnelResponse::error(format!("unknown op: {}", other))),
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Batch handler
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async fn handle_batch(
|
||||
State(state): State<AppState>,
|
||||
body: Bytes,
|
||||
) -> impl IntoResponse {
|
||||
// Decompress if gzipped
|
||||
let json_bytes = if body.starts_with(&[0x1f, 0x8b]) {
|
||||
match decompress_gzip(&body) {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
let resp = serde_json::to_vec(&BatchResponse {
|
||||
r: vec![TunnelResponse::error(format!("gzip decode: {}", e))],
|
||||
}).unwrap_or_default();
|
||||
return (StatusCode::OK, [(header::CONTENT_TYPE, "application/json")], resp);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
body.to_vec()
|
||||
};
|
||||
|
||||
let req: BatchRequest = match serde_json::from_slice(&json_bytes) {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
let resp = serde_json::to_vec(&BatchResponse {
|
||||
r: vec![TunnelResponse::error(format!("bad json: {}", e))],
|
||||
}).unwrap_or_default();
|
||||
return (StatusCode::OK, [(header::CONTENT_TYPE, "application/json")], resp);
|
||||
}
|
||||
};
|
||||
|
||||
if req.k != state.auth_key {
|
||||
let resp = serde_json::to_vec(&BatchResponse {
|
||||
r: vec![TunnelResponse::error("unauthorized")],
|
||||
}).unwrap_or_default();
|
||||
return (StatusCode::OK, [(header::CONTENT_TYPE, "application/json")], resp);
|
||||
}
|
||||
|
||||
// Process all ops. For "data" ops, first write all outbound data,
|
||||
// then do a short sleep to let servers respond, then drain all.
|
||||
// This batches the network round trips on the server side too.
|
||||
|
||||
// Phase 1: process connects and writes
|
||||
let mut results: Vec<(usize, TunnelResponse)> = Vec::with_capacity(req.ops.len());
|
||||
let mut data_ops: Vec<(usize, String)> = Vec::new(); // (index, sid) for data ops needing drain
|
||||
|
||||
for (i, op) in req.ops.iter().enumerate() {
|
||||
match op.op.as_str() {
|
||||
"connect" => {
|
||||
let r = handle_connect(&state, op.host.clone(), op.port).await;
|
||||
results.push((i, r));
|
||||
}
|
||||
"data" => {
|
||||
let sid = match &op.sid {
|
||||
Some(s) if !s.is_empty() => s.clone(),
|
||||
_ => { results.push((i, TunnelResponse::error("missing sid"))); continue; }
|
||||
};
|
||||
|
||||
// Write outbound data
|
||||
let sessions = state.sessions.lock().await;
|
||||
if let Some(session) = sessions.get(&sid) {
|
||||
*session.inner.last_active.lock().await = Instant::now();
|
||||
if let Some(ref data_b64) = op.d {
|
||||
if !data_b64.is_empty() {
|
||||
if let Ok(bytes) = B64.decode(data_b64) {
|
||||
if !bytes.is_empty() {
|
||||
let mut w = session.inner.writer.lock().await;
|
||||
let _ = w.write_all(&bytes).await;
|
||||
let _ = w.flush().await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
drop(sessions);
|
||||
data_ops.push((i, sid));
|
||||
} else {
|
||||
drop(sessions);
|
||||
results.push((i, TunnelResponse { sid: Some(sid), d: None, eof: Some(true), e: None }));
|
||||
}
|
||||
}
|
||||
"close" => {
|
||||
let r = handle_close(&state, op.sid.clone()).await;
|
||||
results.push((i, r));
|
||||
}
|
||||
other => {
|
||||
results.push((i, TunnelResponse::error(format!("unknown op: {}", other))));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 2: short wait for servers to respond, then drain all data sessions
|
||||
if !data_ops.is_empty() {
|
||||
// Give servers a moment to respond to the data we just wrote
|
||||
tokio::time::sleep(Duration::from_millis(150)).await;
|
||||
|
||||
// First drain pass
|
||||
{
|
||||
let sessions = state.sessions.lock().await;
|
||||
let mut need_retry = Vec::new();
|
||||
for (i, sid) in &data_ops {
|
||||
if let Some(session) = sessions.get(sid) {
|
||||
let (data, eof) = drain_now(&session.inner).await;
|
||||
if data.is_empty() && !eof {
|
||||
need_retry.push((*i, sid.clone()));
|
||||
} else {
|
||||
results.push((*i, TunnelResponse {
|
||||
sid: Some(sid.clone()),
|
||||
d: if data.is_empty() { None } else { Some(B64.encode(&data)) },
|
||||
eof: Some(eof), e: None,
|
||||
}));
|
||||
}
|
||||
} else {
|
||||
results.push((*i, TunnelResponse {
|
||||
sid: Some(sid.clone()), d: None, eof: Some(true), e: None,
|
||||
}));
|
||||
}
|
||||
}
|
||||
drop(sessions);
|
||||
|
||||
// Retry sessions that had no data yet
|
||||
if !need_retry.is_empty() {
|
||||
tokio::time::sleep(Duration::from_millis(200)).await;
|
||||
let sessions = state.sessions.lock().await;
|
||||
for (i, sid) in &need_retry {
|
||||
if let Some(s) = sessions.get(sid) {
|
||||
let (data, eof) = drain_now(&s.inner).await;
|
||||
results.push((*i, TunnelResponse {
|
||||
sid: Some(sid.clone()),
|
||||
d: if data.is_empty() { None } else { Some(B64.encode(&data)) },
|
||||
eof: Some(eof), e: None,
|
||||
}));
|
||||
} else {
|
||||
results.push((*i, TunnelResponse {
|
||||
sid: Some(sid.clone()), d: None, eof: Some(true), e: None,
|
||||
}));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Clean up eof sessions
|
||||
let mut sessions = state.sessions.lock().await;
|
||||
for (_, sid) in &data_ops {
|
||||
if let Some(s) = sessions.get(sid) {
|
||||
if s.inner.eof.load(Ordering::Acquire) {
|
||||
if let Some(s) = sessions.remove(sid) {
|
||||
s.reader_handle.abort();
|
||||
tracing::info!("session {} closed by remote (batch)", sid);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Sort results by original index and build response
|
||||
results.sort_by_key(|(i, _)| *i);
|
||||
let batch_resp = BatchResponse {
|
||||
r: results.into_iter().map(|(_, r)| r).collect(),
|
||||
};
|
||||
|
||||
let json = serde_json::to_vec(&batch_resp).unwrap_or_default();
|
||||
(StatusCode::OK, [(header::CONTENT_TYPE, "application/json")], json)
|
||||
}
|
||||
|
||||
fn compress_gzip(data: &[u8]) -> Vec<u8> {
|
||||
use std::io::Write;
|
||||
let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::fast());
|
||||
let _ = encoder.write_all(data);
|
||||
encoder.finish().unwrap_or_else(|_| data.to_vec())
|
||||
}
|
||||
|
||||
fn decompress_gzip(data: &[u8]) -> Result<Vec<u8>, String> {
|
||||
use std::io::Read;
|
||||
let mut decoder = flate2::read::GzDecoder::new(data);
|
||||
let mut out = Vec::new();
|
||||
decoder.read_to_end(&mut out).map_err(|e| e.to_string())?;
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Shared op handlers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async fn handle_connect(state: &AppState, host: Option<String>, port: Option<u16>) -> TunnelResponse {
|
||||
let host = match host {
|
||||
Some(h) if !h.is_empty() => h,
|
||||
_ => return TunnelResponse::error("missing host"),
|
||||
};
|
||||
let port = match port {
|
||||
Some(p) if p > 0 => p,
|
||||
_ => return TunnelResponse::error("missing or invalid port"),
|
||||
};
|
||||
let session = match create_session(&host, port).await {
|
||||
Ok(s) => s,
|
||||
Err(e) => return TunnelResponse::error(format!("connect failed: {}", e)),
|
||||
};
|
||||
let sid = uuid::Uuid::new_v4().to_string();
|
||||
tracing::info!("session {} -> {}:{}", sid, host, port);
|
||||
state.sessions.lock().await.insert(sid.clone(), session);
|
||||
TunnelResponse { sid: Some(sid), d: None, eof: Some(false), e: None }
|
||||
}
|
||||
|
||||
async fn handle_data_single(state: &AppState, sid: Option<String>, data: Option<String>) -> TunnelResponse {
|
||||
let sid = match sid {
|
||||
Some(s) if !s.is_empty() => s,
|
||||
_ => return TunnelResponse::error("missing sid"),
|
||||
};
|
||||
let sessions = state.sessions.lock().await;
|
||||
let session = match sessions.get(&sid) {
|
||||
Some(s) => s,
|
||||
None => return TunnelResponse::error("unknown session"),
|
||||
};
|
||||
*session.inner.last_active.lock().await = Instant::now();
|
||||
if let Some(ref data_b64) = data {
|
||||
if !data_b64.is_empty() {
|
||||
if let Ok(bytes) = B64.decode(data_b64) {
|
||||
if !bytes.is_empty() {
|
||||
let mut w = session.inner.writer.lock().await;
|
||||
if let Err(e) = w.write_all(&bytes).await {
|
||||
drop(w); drop(sessions);
|
||||
state.sessions.lock().await.remove(&sid);
|
||||
return TunnelResponse::error(format!("write failed: {}", e));
|
||||
}
|
||||
let _ = w.flush().await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let (data, eof) = wait_and_drain(&session.inner, Duration::from_secs(5)).await;
|
||||
drop(sessions);
|
||||
if eof {
|
||||
if let Some(s) = state.sessions.lock().await.remove(&sid) {
|
||||
s.reader_handle.abort();
|
||||
tracing::info!("session {} closed by remote", sid);
|
||||
}
|
||||
}
|
||||
TunnelResponse {
|
||||
sid: Some(sid),
|
||||
d: if data.is_empty() { None } else { Some(B64.encode(&data)) },
|
||||
eof: Some(eof), e: None,
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_close(state: &AppState, sid: Option<String>) -> TunnelResponse {
|
||||
let sid = match sid {
|
||||
Some(s) if !s.is_empty() => s,
|
||||
_ => return TunnelResponse::error("missing sid"),
|
||||
};
|
||||
if let Some(s) = state.sessions.lock().await.remove(&sid) {
|
||||
s.reader_handle.abort();
|
||||
tracing::info!("session {} closed by client", sid);
|
||||
}
|
||||
TunnelResponse { sid: Some(sid), d: None, eof: Some(true), e: None }
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Cleanup
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async fn cleanup_task(sessions: Arc<Mutex<HashMap<String, ManagedSession>>>) {
|
||||
let mut interval = tokio::time::interval(Duration::from_secs(30));
|
||||
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||
loop {
|
||||
interval.tick().await;
|
||||
let mut map = sessions.lock().await;
|
||||
let now = Instant::now();
|
||||
let mut stale = Vec::new();
|
||||
for (k, s) in map.iter() {
|
||||
let last = *s.inner.last_active.lock().await;
|
||||
if now.duration_since(last) > Duration::from_secs(300) {
|
||||
stale.push(k.clone());
|
||||
}
|
||||
}
|
||||
for k in &stale {
|
||||
if let Some(s) = map.remove(k) {
|
||||
s.reader_handle.abort();
|
||||
tracing::info!("reaped idle session {}", k);
|
||||
}
|
||||
}
|
||||
if !stale.is_empty() {
|
||||
tracing::info!("cleanup: reaped {}, {} active", stale.len(), map.len());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Main
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| "info".into()),
|
||||
)
|
||||
.init();
|
||||
|
||||
let auth_key = std::env::var("TUNNEL_AUTH_KEY").unwrap_or_else(|_| {
|
||||
tracing::warn!("TUNNEL_AUTH_KEY not set — using default (INSECURE)");
|
||||
"changeme".into()
|
||||
});
|
||||
let port: u16 = std::env::var("PORT")
|
||||
.ok()
|
||||
.and_then(|p| p.parse().ok())
|
||||
.unwrap_or(8080);
|
||||
|
||||
let sessions: Arc<Mutex<HashMap<String, ManagedSession>>> =
|
||||
Arc::new(Mutex::new(HashMap::new()));
|
||||
tokio::spawn(cleanup_task(sessions.clone()));
|
||||
|
||||
let state = AppState { sessions, auth_key };
|
||||
|
||||
let app = Router::new()
|
||||
.route("/tunnel", post(handle_tunnel))
|
||||
.route("/tunnel/batch", post(handle_batch))
|
||||
.route("/health", axum::routing::get(|| async { "ok" }))
|
||||
.with_state(state);
|
||||
|
||||
let addr = format!("0.0.0.0:{}", port);
|
||||
tracing::info!("tunnel-node listening on {}", addr);
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
|
||||
axum::serve(listener, app)
|
||||
.with_graceful_shutdown(async {
|
||||
tokio::signal::ctrl_c().await.ok();
|
||||
tracing::info!("shutting down");
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
Reference in New Issue
Block a user