use std::{
io::{self, IoSlice},
net::SocketAddr,
pin::Pin,
task::{self, Poll},
};
use pin_project::pin_project;
use shadowsocks::net::TcpStream;
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
#[cfg(feature = "local-http")]
use super::http_connect::HttpConnectTunnel;
#[cfg(any(feature = "local-http-native-tls", feature = "local-http-rustls"))]
use super::tls::OutboundTlsStream;
pub struct OutboundProxyStream {
local_addr: SocketAddr,
inner: OutboundProxyStreamInner,
}
#[allow(clippy::large_enum_variant)]
#[pin_project(project = OutboundProxyStreamInnerProj)]
enum OutboundProxyStreamInner {
Bypassed(#[pin] TcpStream),
#[cfg(any(feature = "local-http-native-tls", feature = "local-http-rustls"))]
Https(#[pin] OutboundTlsStream),
#[cfg(feature = "local-http")]
Http(#[pin] HttpConnectTunnel),
}
impl OutboundProxyStream {
pub fn from_tcp(stream: TcpStream) -> io::Result<Self> {
let local_addr = stream.local_addr()?;
Ok(Self {
local_addr,
inner: OutboundProxyStreamInner::Bypassed(stream),
})
}
pub fn from_tcp_with_local_addr(stream: TcpStream, local_addr: SocketAddr) -> Self {
Self {
local_addr,
inner: OutboundProxyStreamInner::Bypassed(stream),
}
}
pub fn local_addr(&self) -> io::Result<SocketAddr> {
Ok(self.local_addr)
}
#[allow(clippy::result_large_err)]
pub fn try_into_tcp(self) -> Result<TcpStream, Self> {
match self.inner {
OutboundProxyStreamInner::Bypassed(s) => Ok(s),
other => Err(Self {
local_addr: self.local_addr,
inner: other,
}),
}
}
#[cfg(any(feature = "local-http-native-tls", feature = "local-http-rustls"))]
pub(super) fn from_tls(local_addr: SocketAddr, tls: OutboundTlsStream) -> Self {
Self {
local_addr,
inner: OutboundProxyStreamInner::Https(tls),
}
}
#[cfg(feature = "local-http")]
pub(super) fn from_http(local_addr: SocketAddr, tunnel: HttpConnectTunnel) -> Self {
Self {
local_addr,
inner: OutboundProxyStreamInner::Http(tunnel),
}
}
fn project_inner(self: Pin<&mut Self>) -> OutboundProxyStreamInnerProj<'_> {
unsafe {
let this = self.get_unchecked_mut();
Pin::new_unchecked(&mut this.inner).project()
}
}
}
impl Unpin for OutboundProxyStream {}
impl AsyncRead for OutboundProxyStream {
fn poll_read(self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
match self.project_inner() {
OutboundProxyStreamInnerProj::Bypassed(s) => s.poll_read(cx, buf),
#[cfg(any(feature = "local-http-native-tls", feature = "local-http-rustls"))]
OutboundProxyStreamInnerProj::Https(s) => s.poll_read(cx, buf),
#[cfg(feature = "local-http")]
OutboundProxyStreamInnerProj::Http(s) => s.poll_read(cx, buf),
}
}
}
impl AsyncWrite for OutboundProxyStream {
fn poll_write(self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
match self.project_inner() {
OutboundProxyStreamInnerProj::Bypassed(s) => s.poll_write(cx, buf),
#[cfg(any(feature = "local-http-native-tls", feature = "local-http-rustls"))]
OutboundProxyStreamInnerProj::Https(s) => s.poll_write(cx, buf),
#[cfg(feature = "local-http")]
OutboundProxyStreamInnerProj::Http(s) => s.poll_write(cx, buf),
}
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
match self.project_inner() {
OutboundProxyStreamInnerProj::Bypassed(s) => s.poll_flush(cx),
#[cfg(any(feature = "local-http-native-tls", feature = "local-http-rustls"))]
OutboundProxyStreamInnerProj::Https(s) => s.poll_flush(cx),
#[cfg(feature = "local-http")]
OutboundProxyStreamInnerProj::Http(s) => s.poll_flush(cx),
}
}
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
match self.project_inner() {
OutboundProxyStreamInnerProj::Bypassed(s) => s.poll_shutdown(cx),
#[cfg(any(feature = "local-http-native-tls", feature = "local-http-rustls"))]
OutboundProxyStreamInnerProj::Https(s) => s.poll_shutdown(cx),
#[cfg(feature = "local-http")]
OutboundProxyStreamInnerProj::Http(s) => s.poll_shutdown(cx),
}
}
fn poll_write_vectored(
self: Pin<&mut Self>,
cx: &mut task::Context<'_>,
bufs: &[IoSlice<'_>],
) -> Poll<io::Result<usize>> {
match self.project_inner() {
OutboundProxyStreamInnerProj::Bypassed(s) => s.poll_write_vectored(cx, bufs),
#[cfg(any(feature = "local-http-native-tls", feature = "local-http-rustls"))]
OutboundProxyStreamInnerProj::Https(s) => s.poll_write_vectored(cx, bufs),
#[cfg(feature = "local-http")]
OutboundProxyStreamInnerProj::Http(s) => s.poll_write_vectored(cx, bufs),
}
}
fn is_write_vectored(&self) -> bool {
match &self.inner {
OutboundProxyStreamInner::Bypassed(s) => s.is_write_vectored(),
#[cfg(any(feature = "local-http-native-tls", feature = "local-http-rustls"))]
OutboundProxyStreamInner::Https(_) => false,
#[cfg(feature = "local-http")]
OutboundProxyStreamInner::Http(_) => false,
}
}
}