use std::sync::Arc;
use async_trait::async_trait;
use pingora::{
apps::ServerApp,
protocols::{Stream, ALPN},
server::ShutdownWatch,
};
use tokio::net::TcpStream;
use crate::proxy_server::cert_store::ServerName;
use super::{choose, TcpTarget};
pub struct TcpProxy {
targets: Arc<Vec<TcpTarget>>,
}
impl TcpProxy {
pub fn new(targets: Arc<Vec<TcpTarget>>) -> Self {
Self { targets }
}
pub fn target(&self, downstream: &Stream) -> Option<&TcpTarget> {
let local = downstream
.get_socket_digest()
.and_then(|digest| digest.local_addr().and_then(|addr| addr.as_inet()).copied())?;
let sni = downstream.get_ssl_digest().and_then(|digest| {
digest
.extension
.get::<ServerName>()
.map(|name| name.0.clone())
});
let alpn = match downstream.selected_alpn_proto() {
Some(ALPN::Custom(protocol)) => Some(protocol),
_ => None,
};
let target = choose(
&self.targets,
&local,
sni.as_deref(),
alpn.as_ref().map(pingora::protocols::tls::CustomALPN::protocol),
);
if target.is_none() {
tracing::debug!("no tcp block for a connection on {local} (server name: {sni:?})");
}
target
}
pub async fn forward(&self, mut downstream: Stream, target: &TcpTarget) {
let Some(addr) = target
.load_balancer
.select(b"", 32)
.and_then(|backend| backend.addr.as_inet().copied())
else {
tracing::warn!("{}: no healthy upstream", target.label);
return;
};
let mut upstream = match TcpStream::connect(addr).await {
Ok(upstream) => upstream,
Err(err) => {
tracing::warn!(
"{}: could not connect to upstream {addr}: {err}",
target.label
);
return;
}
};
if let Err(err) = tokio::io::copy_bidirectional(&mut downstream, &mut upstream).await {
tracing::debug!("{}: connection to {addr} ended: {err}", target.label);
}
}
}
#[async_trait]
impl ServerApp for TcpProxy {
async fn process_new(
self: &Arc<Self>,
downstream: Stream,
_shutdown: &ShutdownWatch,
) -> Option<Stream> {
let target = self.target(&downstream)?;
self.forward(downstream, target).await;
None
}
}