use std::collections::BTreeMap;
use std::convert::Infallible;
use std::fs::OpenOptions;
use std::future::Future;
use std::io::{self, BufWriter, Write};
use std::net::SocketAddr;
use std::path::Path;
use std::pin::Pin;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, mpsc};
use std::task::{Context, Poll};
use std::thread::JoinHandle;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use bytes::Bytes;
use http_body_util::combinators::BoxBody;
use http_body_util::{BodyExt, Empty, Full};
use hyper::body::{Body, Frame, Incoming, SizeHint};
use hyper::client::conn::http1 as client_http1;
use hyper::header::{self, HeaderMap, HeaderName};
use hyper::server::conn::http1 as server_http1;
use hyper::service::service_fn;
use hyper::{Request, Response, StatusCode, Uri, Version};
use hyper_util::rt::TokioIo;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{oneshot, watch};
use crate::RECORD_VERSION;
use crate::normalize;
use crate::record::{
self, ExchangeError, Record, RequestRecord, ResponseRecord, StreamKind, write_jsonl,
};
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
const JSONL_QUEUE: usize = 1024;
const RECORDED_HEADERS: [&str; 4] = [
"user-agent",
"content-type",
"content-length",
"transfer-encoding",
];
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ProxyConfig {
pub listen: SocketAddr,
pub upstream: SocketAddr,
pub body_mode: BodyMode,
pub max_body_bytes: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum BodyMode {
#[default]
Full,
Digest,
}
pub trait Sink: Send + Sync + 'static {
fn record(&self, record: Record);
}
impl<F: Fn(Record) + Send + Sync + 'static> Sink for F {
fn record(&self, record: Record) {
self(record)
}
}
impl Sink for mpsc::Sender<Record> {
fn record(&self, record: Record) {
let _ = self.send(record);
}
}
#[derive(Debug)]
pub struct JsonlSink {
sender: Option<mpsc::SyncSender<Record>>,
writer: Option<JoinHandle<io::Result<u64>>>,
}
impl JsonlSink {
pub fn create(path: &Path) -> io::Result<JsonlSink> {
let file = OpenOptions::new().create(true).append(true).open(path)?;
let (sender, receiver) = mpsc::sync_channel(JSONL_QUEUE);
let writer = std::thread::Builder::new()
.name("llmp-jsonl".into())
.spawn(move || write_records(file, receiver))?;
Ok(JsonlSink {
sender: Some(sender),
writer: Some(writer),
})
}
pub fn close(mut self) -> io::Result<u64> {
self.shut_down()
}
fn shut_down(&mut self) -> io::Result<u64> {
self.sender = None;
match self.writer.take() {
Some(writer) => writer
.join()
.unwrap_or_else(|_| Err(io::Error::other("the JSONL writer thread panicked"))),
None => Ok(0),
}
}
}
impl Sink for JsonlSink {
fn record(&self, record: Record) {
if let Some(sender) = &self.sender {
let _ = sender.send(record);
}
}
}
impl Drop for JsonlSink {
fn drop(&mut self) {
let _ = self.shut_down();
}
}
fn write_records(file: std::fs::File, receiver: mpsc::Receiver<Record>) -> io::Result<u64> {
let mut out = BufWriter::new(file);
let mut written = 0u64;
let mut first_error = None;
for record in receiver {
if first_error.is_some() {
continue;
}
match write_jsonl(&mut out, &record).and_then(|()| out.flush()) {
Ok(()) => written += 1,
Err(e) => first_error = Some(e),
}
}
match first_error {
Some(e) => Err(e),
None => Ok(written),
}
}
pub async fn serve(
config: ProxyConfig,
sink: impl Sink,
shutdown: impl Future<Output = ()>,
) -> io::Result<()> {
let listener = TcpListener::bind(config.listen).await?;
let shared = Arc::new(Shared {
upstream: config.upstream,
body_mode: config.body_mode,
max_body_bytes: config.max_body_bytes,
sink: Arc::new(sink),
next_id: AtomicU64::new(1),
});
let (stop_tx, stop_rx) = watch::channel(false);
let (alive_tx, mut alive_rx) = tokio::sync::mpsc::channel::<()>(1);
let mut next_conn = 1u64;
tokio::pin!(shutdown);
loop {
tokio::select! {
() = &mut shutdown => break,
accepted = listener.accept() => match accepted {
Ok((stream, _)) => {
let _ = stream.set_nodelay(true);
let connection = Arc::new(Connection {
id: next_conn,
shared: Arc::clone(&shared),
upstream: tokio::sync::Mutex::new(None),
delivered: Mutex::new(None),
alive: alive_tx.clone(),
});
next_conn += 1;
tokio::spawn(run_connection(stream, connection, stop_rx.clone()));
}
Err(_) => tokio::time::sleep(Duration::from_millis(50)).await,
},
}
}
drop(listener);
let _ = stop_tx.send(true);
drop(alive_tx);
let _ = alive_rx.recv().await;
Ok(())
}
pub fn serve_blocking(config: ProxyConfig, sink: impl Sink) -> io::Result<()> {
let runtime = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()?;
runtime.block_on(serve(config, sink, async {
if tokio::signal::ctrl_c().await.is_err() {
std::future::pending::<()>().await;
}
}))
}
struct Shared {
upstream: SocketAddr,
body_mode: BodyMode,
max_body_bytes: usize,
sink: Arc<dyn Sink>,
next_id: AtomicU64,
}
struct Connection {
id: u64,
shared: Arc<Shared>,
upstream: tokio::sync::Mutex<Option<client_http1::SendRequest<Full<Bytes>>>>,
delivered: Mutex<Option<oneshot::Receiver<()>>>,
alive: tokio::sync::mpsc::Sender<()>,
}
type ProxyBody = BoxBody<Bytes, hyper::Error>;
async fn run_connection(
stream: TcpStream,
connection: Arc<Connection>,
mut stop: watch::Receiver<bool>,
) {
let service = service_fn(move |request| {
let connection = Arc::clone(&connection);
async move { Ok::<_, Infallible>(forward(connection, request).await) }
});
let conn = server_http1::Builder::new().serve_connection(TokioIo::new(stream), service);
tokio::pin!(conn);
let stopped = tokio::select! {
_ = conn.as_mut() => false,
result = stop.wait_for(|stopped| *stopped) => {
drop(result);
true
}
};
if stopped {
conn.as_mut().graceful_shutdown();
let _ = conn.await;
}
}
async fn forward(connection: Arc<Connection>, request: Request<Incoming>) -> Response<ProxyBody> {
let t_start_ms = now_ms();
let shared = Arc::clone(&connection.shared);
let id = shared.next_id.fetch_add(1, Ordering::Relaxed);
let (delivered_tx, delivered_rx) = oneshot::channel();
let previous = connection
.delivered
.lock()
.ok()
.and_then(|mut slot| slot.replace(delivered_rx));
let (mut parts, body) = request.into_parts();
let path = parts.uri.path_and_query().map_or_else(
|| parts.uri.path().to_string(),
|pq| pq.as_str().to_string(),
);
let mut flight = InFlight {
capture: Some(Capture {
id,
conn: connection.id,
t_start_ms,
method: parts.method.to_string(),
path,
request_headers: recorded_headers(&parts.headers),
request_body: Bytes::new(),
request_body_bytes: 0,
request_too_large: false,
response: None,
t_first_byte_ms: None,
}),
delivery: Some(Delivery {
shared: Arc::clone(&shared),
delivered: delivered_tx,
alive: connection.alive.clone(),
}),
};
if let Some(previous) = previous {
let _ = previous.await;
}
let body = match body.collect().await {
Ok(collected) => collected.to_bytes(),
Err(_) => {
flight.finish(Some(ExchangeError::ClientDisconnected));
return status_response(StatusCode::BAD_REQUEST);
}
};
if let Some(capture) = flight.capture.as_mut() {
capture.request_body_bytes = body.len() as u64;
capture.request_too_large = body.len() > shared.max_body_bytes;
if !capture.request_too_large {
capture.request_body = body.clone();
}
}
strip_hop_by_hop(&mut parts.headers);
let upstream_request = UpstreamRequest {
method: parts.method,
uri: origin_form(parts.uri),
headers: parts.headers,
body,
};
let response = match connection.send(&upstream_request).await {
Ok(response) => response,
Err(error) => {
flight.finish(Some(error));
return status_response(StatusCode::BAD_GATEWAY);
}
};
let (mut parts, incoming) = response.into_parts();
if let Some(capture) = flight.capture.as_mut() {
capture.response = Some(ResponseCapture {
status: parts.status.as_u16(),
headers: recorded_headers(&parts.headers),
stream: stream_kind(&parts.headers),
body: Vec::new(),
body_bytes: 0,
truncated: false,
});
}
strip_hop_by_hop(&mut parts.headers);
let body = RecordingBody {
inner: incoming,
flight,
max_body_bytes: shared.max_body_bytes,
};
Response::from_parts(parts, body.boxed())
}
impl Connection {
async fn send(&self, request: &UpstreamRequest) -> Result<Response<Incoming>, ExchangeError> {
let mut slot = self.upstream.lock().await;
for _ in 0..2 {
let reused = slot.as_ref().is_some_and(|sender| !sender.is_closed());
if !reused {
*slot = Some(connect(self.shared.upstream).await?);
}
let Some(sender) = slot.as_mut() else {
return Err(ExchangeError::UpstreamUnreachable);
};
if sender.ready().await.is_err() {
*slot = None;
if reused {
continue;
}
return Err(ExchangeError::UpstreamReset);
}
match sender.send_request(request.build()).await {
Ok(response) => return Ok(response),
Err(e) => {
*slot = None;
if reused && e.is_canceled() {
continue;
}
return Err(ExchangeError::UpstreamReset);
}
}
}
Err(ExchangeError::UpstreamReset)
}
}
async fn connect(
addr: SocketAddr,
) -> Result<client_http1::SendRequest<Full<Bytes>>, ExchangeError> {
let stream = match tokio::time::timeout(CONNECT_TIMEOUT, TcpStream::connect(addr)).await {
Ok(Ok(stream)) => stream,
Ok(Err(_)) => return Err(ExchangeError::UpstreamUnreachable),
Err(_) => return Err(ExchangeError::Timeout),
};
let _ = stream.set_nodelay(true);
let (sender, conn) = client_http1::handshake(TokioIo::new(stream))
.await
.map_err(|_| ExchangeError::UpstreamUnreachable)?;
tokio::spawn(async move {
let _ = conn.await;
});
Ok(sender)
}
struct UpstreamRequest {
method: hyper::Method,
uri: Uri,
headers: HeaderMap,
body: Bytes,
}
impl UpstreamRequest {
fn build(&self) -> Request<Full<Bytes>> {
let mut request = Request::new(Full::new(self.body.clone()));
*request.method_mut() = self.method.clone();
*request.uri_mut() = self.uri.clone();
*request.headers_mut() = self.headers.clone();
*request.version_mut() = Version::HTTP_11;
request
}
}
struct Capture {
id: u64,
conn: u64,
t_start_ms: i64,
method: String,
path: String,
request_headers: BTreeMap<String, String>,
request_body: Bytes,
request_body_bytes: u64,
request_too_large: bool,
response: Option<ResponseCapture>,
t_first_byte_ms: Option<i64>,
}
struct ResponseCapture {
status: u16,
headers: BTreeMap<String, String>,
stream: Option<StreamKind>,
body: Vec<u8>,
body_bytes: u64,
truncated: bool,
}
struct Delivery {
shared: Arc<Shared>,
delivered: oneshot::Sender<()>,
alive: tokio::sync::mpsc::Sender<()>,
}
impl Delivery {
fn deliver(self, capture: Capture, error: Option<ExchangeError>) {
let t_end_ms = now_ms();
let Delivery {
shared,
delivered,
alive,
} = self;
let job = move || {
let record = build_record(capture, error, t_end_ms, shared.body_mode);
shared.sink.record(record);
let _ = delivered.send(());
drop(alive);
};
match tokio::runtime::Handle::try_current() {
Ok(handle) => {
handle.spawn_blocking(job);
}
Err(_) => job(),
}
}
}
struct InFlight {
capture: Option<Capture>,
delivery: Option<Delivery>,
}
impl InFlight {
fn finish(&mut self, error: Option<ExchangeError>) {
if let (Some(capture), Some(delivery)) = (self.capture.take(), self.delivery.take()) {
delivery.deliver(capture, error);
}
}
}
impl Drop for InFlight {
fn drop(&mut self) {
self.finish(Some(ExchangeError::ClientDisconnected));
}
}
struct RecordingBody {
inner: Incoming,
flight: InFlight,
max_body_bytes: usize,
}
impl RecordingBody {
fn on_data(&mut self, data: &Bytes) {
let Some(capture) = self.flight.capture.as_mut() else {
return;
};
if capture.t_first_byte_ms.is_none() && !data.is_empty() {
capture.t_first_byte_ms = Some(now_ms());
}
let Some(response) = capture.response.as_mut() else {
return;
};
response.body_bytes += data.len() as u64;
let room = self.max_body_bytes.saturating_sub(response.body.len());
if data.len() > room {
response.truncated = true;
}
response
.body
.extend_from_slice(&data[..data.len().min(room)]);
}
}
impl Body for RecordingBody {
type Data = Bytes;
type Error = hyper::Error;
fn poll_frame(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Bytes>, hyper::Error>>> {
let this = self.get_mut();
match Pin::new(&mut this.inner).poll_frame(cx) {
Poll::Ready(Some(Ok(frame))) => {
if let Some(data) = frame.data_ref() {
this.on_data(data);
}
Poll::Ready(Some(Ok(frame)))
}
Poll::Ready(Some(Err(e))) => {
this.flight.finish(Some(ExchangeError::UpstreamReset));
Poll::Ready(Some(Err(e)))
}
Poll::Ready(None) => {
this.flight.finish(None);
Poll::Ready(None)
}
Poll::Pending => Poll::Pending,
}
}
fn is_end_stream(&self) -> bool {
self.inner.is_end_stream()
}
fn size_hint(&self) -> SizeHint {
self.inner.size_hint()
}
}
impl Drop for RecordingBody {
fn drop(&mut self) {
let error = if self.inner.is_end_stream() {
None
} else {
Some(ExchangeError::ClientDisconnected)
};
self.flight.finish(error);
}
}
fn build_record(
capture: Capture,
error: Option<ExchangeError>,
t_end_ms: i64,
body_mode: BodyMode,
) -> Record {
let (body, body_text) = if capture.request_too_large {
(None, None)
} else {
stored_body(&capture.request_body)
};
let error = error.or(capture
.request_too_large
.then_some(ExchangeError::BodyTooLarge));
let response = capture.response.map(|response| {
let (body, body_text, chunks, bad_chunks) = match response.stream {
Some(kind) => {
let merged = normalize::merge_stream(kind, &response.body);
let (body, body_text) = if merged.body.is_null() {
(None, lossy_text(&response.body))
} else {
(Some(merged.body), None)
};
(
body,
body_text,
Some(merged.chunks),
Some(merged.bad_chunks),
)
}
None => {
let (body, body_text) = stored_body(&response.body);
(body, body_text, None, None)
}
};
ResponseRecord {
status: response.status,
headers: response.headers,
stream: response.stream,
chunks,
bad_chunks,
body_bytes: response.body_bytes,
truncated: response.truncated,
body,
body_text,
digest: None,
}
});
let mut record = Record {
v: RECORD_VERSION,
id: capture.id,
conn: capture.conn,
t_start_ms: capture.t_start_ms,
t_first_byte_ms: capture.t_first_byte_ms,
t_end_ms,
request: RequestRecord {
method: capture.method,
path: capture.path,
headers: capture.request_headers,
body_bytes: capture.request_body_bytes,
body,
body_text,
digest: None,
},
response,
error,
tokens: None,
};
if body_mode == BodyMode::Digest {
record::apply_digest(&mut record);
}
record
}
fn stored_body(bytes: &[u8]) -> (Option<Value>, Option<String>) {
match serde_json::from_slice::<Value>(bytes) {
Ok(value) if !value.is_null() => (Some(value), None),
_ => (None, lossy_text(bytes)),
}
}
fn lossy_text(bytes: &[u8]) -> Option<String> {
(!bytes.is_empty()).then(|| String::from_utf8_lossy(bytes).into_owned())
}
fn recorded_headers(headers: &HeaderMap) -> BTreeMap<String, String> {
RECORDED_HEADERS
.iter()
.filter_map(|name| {
headers.get(*name).map(|value| {
(
name.to_string(),
String::from_utf8_lossy(value.as_bytes()).into_owned(),
)
})
})
.collect()
}
fn stream_kind(headers: &HeaderMap) -> Option<StreamKind> {
let content_type = headers
.get(header::CONTENT_TYPE)?
.to_str()
.ok()?
.to_ascii_lowercase();
if content_type.contains("text/event-stream") {
Some(StreamKind::Sse)
} else if content_type.contains("ndjson") {
Some(StreamKind::Ndjson)
} else {
None
}
}
fn strip_hop_by_hop(headers: &mut HeaderMap) {
let listed: Vec<HeaderName> = headers
.get_all(header::CONNECTION)
.iter()
.filter_map(|value| value.to_str().ok())
.flat_map(|value| value.split(','))
.filter_map(|name| HeaderName::from_bytes(name.trim().as_bytes()).ok())
.collect();
for name in listed {
headers.remove(name);
}
for name in [
header::CONNECTION,
HeaderName::from_static("proxy-connection"),
HeaderName::from_static("keep-alive"),
header::TE,
header::TRANSFER_ENCODING,
header::UPGRADE,
] {
headers.remove(name);
}
}
fn origin_form(uri: Uri) -> Uri {
if uri.scheme().is_none() {
return uri;
}
uri.path_and_query()
.and_then(|pq| pq.as_str().parse().ok())
.unwrap_or_else(|| Uri::from_static("/"))
}
fn status_response(status: StatusCode) -> Response<ProxyBody> {
let mut response = Response::new(Empty::new().map_err(|never| match never {}).boxed());
*response.status_mut() = status;
response
}
fn now_ms() -> i64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_or(0, |d| i64::try_from(d.as_millis()).unwrap_or(i64::MAX))
}
#[cfg(test)]
mod tests {
use super::*;
use hyper::header::HeaderValue;
#[test]
fn hop_by_hop_headers_are_removed() {
let mut headers = HeaderMap::new();
headers.insert(
header::CONNECTION,
HeaderValue::from_static("keep-alive, X-Private"),
);
headers.insert("x-private", HeaderValue::from_static("1"));
headers.insert("keep-alive", HeaderValue::from_static("timeout=5"));
headers.insert("proxy-connection", HeaderValue::from_static("keep-alive"));
headers.insert(header::TE, HeaderValue::from_static("trailers"));
headers.insert(
header::TRANSFER_ENCODING,
HeaderValue::from_static("chunked"),
);
headers.insert(header::UPGRADE, HeaderValue::from_static("h2c"));
headers.insert(
header::CONTENT_TYPE,
HeaderValue::from_static("application/json"),
);
headers.insert(header::AUTHORIZATION, HeaderValue::from_static("Bearer x"));
strip_hop_by_hop(&mut headers);
let mut left: Vec<&str> = headers.keys().map(HeaderName::as_str).collect();
left.sort();
assert_eq!(left, vec!["authorization", "content-type"]);
}
#[test]
fn only_listed_headers_are_recorded() {
let mut headers = HeaderMap::new();
headers.insert(header::AUTHORIZATION, HeaderValue::from_static("Bearer x"));
headers.insert(header::USER_AGENT, HeaderValue::from_static("t/1"));
headers.insert(header::CONTENT_LENGTH, HeaderValue::from_static("3"));
headers.insert(header::COOKIE, HeaderValue::from_static("a=b"));
let recorded = recorded_headers(&headers);
assert_eq!(
recorded.keys().collect::<Vec<_>>(),
vec!["content-length", "user-agent"]
);
}
#[test]
fn stream_kind_from_content_type() {
let mut headers = HeaderMap::new();
assert_eq!(stream_kind(&headers), None);
headers.insert(
header::CONTENT_TYPE,
HeaderValue::from_static("text/event-stream; charset=utf-8"),
);
assert_eq!(stream_kind(&headers), Some(StreamKind::Sse));
headers.insert(
header::CONTENT_TYPE,
HeaderValue::from_static("application/x-ndjson"),
);
assert_eq!(stream_kind(&headers), Some(StreamKind::Ndjson));
headers.insert(
header::CONTENT_TYPE,
HeaderValue::from_static("application/json"),
);
assert_eq!(stream_kind(&headers), None);
}
#[test]
fn absolute_uris_become_origin_form() {
assert_eq!(
origin_form(Uri::from_static("http://h:1/v1/x?a=1")),
"/v1/x?a=1"
);
assert_eq!(origin_form(Uri::from_static("/v1/x")), "/v1/x");
}
#[test]
fn stored_body_rules() {
assert_eq!(stored_body(b""), (None, None));
assert_eq!(
stored_body(b"{\"a\":1}"),
(Some(serde_json::json!({"a": 1})), None)
);
assert_eq!(stored_body(b"null"), (None, Some("null".into())));
assert_eq!(
stored_body(b"oops \xff"),
(None, Some("oops \u{fffd}".into()))
);
}
}