use rama_core::{
Service,
error::{BoxError, ErrorContext as _},
extensions,
io::{BridgeIo, Io},
rt::Executor,
};
use rama_net::proxy::IoForwardService;
use super::Socks5MitmHandshakeOutcome;
#[derive(Debug, Clone)]
pub struct Socks5MitmRelayService<I, F = IoForwardService> {
dpi_svc: I,
fallback_svc: F,
}
impl<I> Socks5MitmRelayService<I> {
pub fn new(exec: Executor, dpi_svc: I) -> Self {
Self {
dpi_svc,
fallback_svc: IoForwardService::new(exec),
}
}
pub fn with_fallback<F>(self, fallback_svc: F) -> Socks5MitmRelayService<I, F> {
Socks5MitmRelayService {
dpi_svc: self.dpi_svc,
fallback_svc,
}
}
}
impl<I, F, Ingress, Egress> Service<BridgeIo<Ingress, Egress>> for Socks5MitmRelayService<I, F>
where
I: Service<BridgeIo<Ingress, Egress>, Error: Into<BoxError>>,
F: Service<BridgeIo<Ingress, Egress>, Error: Into<BoxError>>,
Ingress: Io + Unpin + extensions::ExtensionsRef,
Egress: Io + Unpin + extensions::ExtensionsRef,
{
type Output = ();
type Error = BoxError;
async fn serve(
&self,
BridgeIo(mut ingress_stream, mut egress_stream): BridgeIo<Ingress, Egress>,
) -> Result<Self::Output, Self::Error> {
let outcome = super::socks5_mitm_relay_handshake(&mut ingress_stream, &mut egress_stream)
.await
.context("socks5 relay handshake using provided I/O bridge")?;
match outcome {
Socks5MitmHandshakeOutcome::ContinueInspection => self
.dpi_svc
.serve(BridgeIo(ingress_stream, egress_stream))
.await
.map(drop)
.context("serve socks5 handshake-relayed bridge I/O using DPI svc"),
Socks5MitmHandshakeOutcome::UnsupportedFlow => self
.fallback_svc
.serve(BridgeIo(ingress_stream, egress_stream))
.await
.map(drop)
.context("serve socks5 handshake-relayed bridge I/O using fallback svc"),
}
}
}