use tokio::sync::Mutex;
use crate::connection::ProxyConnection;
use crate::stream::reader::ProxyReader;
use crate::stream::writer::ProxyWriter;
use crate::{ErrorKind, Proxy, ProxyError, ProxyProtocol, ProxyResult};
pub struct ProxyStream {
pub reader: Mutex<Option<ProxyReader>>,
pub writer: Mutex<Option<ProxyWriter>>,
proxy: Proxy,
}
impl ProxyStream {
pub fn new(proxy: impl Into<Proxy>) -> Self {
Self {
proxy: proxy.into(),
reader: Mutex::new(None),
writer: Mutex::new(None),
}
}
pub fn set_proxy(&mut self, proxy: impl Into<Proxy>) {
self.proxy = proxy.into();
}
#[deprecated = "method was renamed, use `.proxy_protocol()`"]
pub fn get_proxy_protocol(&self) -> &ProxyProtocol {
self.proxy.protocol()
}
pub fn proxy_protocol(&self) -> &ProxyProtocol {
self.proxy.protocol()
}
pub async fn connect(&self, target_host: impl Into<String>, target_port: u16) -> ProxyResult<()> {
let stream = self.proxy.connect(target_host, target_port).await?;
match stream {
#[cfg(any(feature = "http", feature = "socks4", feature = "socks5"))]
ProxyConnection::Tcp(tcp_stream) => {
let (rh, wh) = tcp_stream.into_split();
*self.reader.lock().await = Some(ProxyReader::new(rh));
*self.writer.lock().await = Some(ProxyWriter::new(wh));
}
#[cfg(feature = "https")]
ProxyConnection::Tls(tls_stream) => {
let (rh, wh) = tokio::io::split(tls_stream);
*self.reader.lock().await = Some(ProxyReader::new(rh));
*self.writer.lock().await = Some(ProxyWriter::new(wh));
}
}
Ok(())
}
pub async fn shutdown(&self) -> ProxyResult<()> {
if let Some(writer) = self.writer.lock().await.as_mut() {
writer.shutdown().await?;
}
*self.reader.lock().await = None;
*self.writer.lock().await = None;
Ok(())
}
pub async fn read(&self, buffer: &mut [u8]) -> ProxyResult<usize> {
if let Some(reader) = self.reader.lock().await.as_mut() {
Ok(reader.read(buffer).await?)
} else {
Err(ProxyError::new(ErrorKind::StreamError, "reader is not initialized"))
}
}
pub async fn read_to_end(&self, buffer: &mut Vec<u8>) -> ProxyResult<usize> {
if let Some(reader) = self.reader.lock().await.as_mut() {
Ok(reader.read_to_end(buffer).await?)
} else {
Err(ProxyError::new(ErrorKind::StreamError, "reader is not initialized"))
}
}
pub async fn write(&self, buffer: &[u8]) -> ProxyResult<()> {
if let Some(writer) = self.writer.lock().await.as_mut() {
Ok(writer.write(buffer).await?)
} else {
Err(ProxyError::new(ErrorKind::StreamError, "writer is not initialized"))
}
}
}
impl From<String> for ProxyStream {
fn from(value: String) -> Self {
Self::new(value)
}
}
impl From<&str> for ProxyStream {
fn from(value: &str) -> Self {
Self::new(value)
}
}