69 lines
2.1 KiB
Rust
69 lines
2.1 KiB
Rust
|
use crate::servers::upstream_address::UpstreamAddress;
|
||
|
|
|
||
|
|
use crate::upstreams::copy;
|
||
|
|
use futures::future::try_join;
|
||
|
|
use log::{debug, error};
|
||
|
|
use serde::Deserialize;
|
||
|
|
use std::net::SocketAddr;
|
||
|
|
use tokio::io;
|
||
|
|
use tokio::net::TcpStream;
|
||
|
|
use tokio::sync::Mutex;
|
||
|
|
|
||
|
|
#[derive(Debug, Default)]
|
||
|
|
struct Addr(Mutex<UpstreamAddress>);
|
||
|
|
|
||
|
|
impl Clone for Addr {
|
||
|
|
fn clone(&self) -> Self {
|
||
|
|
tokio::task::block_in_place(|| Self(Mutex::new(self.0.blocking_lock().clone())))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
#[derive(Debug, Clone, Deserialize, Default)]
|
||
|
|
pub struct ProxyToUpstream {
|
||
|
|
pub addr: String,
|
||
|
|
pub protocol: String,
|
||
|
|
#[serde(skip_deserializing)]
|
||
|
|
addresses: Addr,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl ProxyToUpstream {
|
||
|
|
pub async fn resolve_addresses(&self) -> std::io::Result<Vec<SocketAddr>> {
|
||
|
|
let mut addr = self.addresses.0.lock().await;
|
||
|
|
addr.resolve((*self.protocol).into()).await
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn new(address: String, protocol: String) -> Self {
|
||
|
|
Self {
|
||
|
|
addr: address.clone(),
|
||
|
|
protocol,
|
||
|
|
addresses: Addr(Mutex::new(UpstreamAddress::new(address))),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
pub(crate) async fn proxy(&self, inbound: TcpStream) -> Result<(), Box<dyn std::error::Error>> {
|
||
|
|
let outbound = match self.protocol.as_ref() {
|
||
|
|
"tcp4" | "tcp6" | "tcp" => {
|
||
|
|
TcpStream::connect(self.resolve_addresses().await?.as_slice()).await?
|
||
|
|
}
|
||
|
|
_ => {
|
||
|
|
error!("Reached unknown protocol: {:?}", self.protocol);
|
||
|
|
return Err("Reached unknown protocol".into());
|
||
|
|
}
|
||
|
|
};
|
||
|
|
|
||
|
|
debug!("Connected to {:?}", outbound.peer_addr().unwrap());
|
||
|
|
|
||
|
|
let (mut ri, mut wi) = io::split(inbound);
|
||
|
|
let (mut ro, mut wo) = io::split(outbound);
|
||
|
|
|
||
|
|
let inbound_to_outbound = copy(&mut ri, &mut wo);
|
||
|
|
let outbound_to_inbound = copy(&mut ro, &mut wi);
|
||
|
|
|
||
|
|
let (bytes_tx, bytes_rx) = try_join(inbound_to_outbound, outbound_to_inbound).await?;
|
||
|
|
|
||
|
|
debug!("Bytes read: {:?} write: {:?}", bytes_tx, bytes_rx);
|
||
|
|
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
}
|