relay_core_lib/capture/
source.rs1use std::net::SocketAddr;
2use std::future::Future;
3use std::pin::Pin;
4use tokio::io::{AsyncRead, AsyncWrite};
5
6pub struct IncomingConnection<IO> {
8 pub stream: IO,
9 pub client_addr: SocketAddr,
10 pub target_addr: Option<SocketAddr>,
12}
13
14pub trait CaptureSource {
18 type IO: AsyncRead + AsyncWrite + Unpin + Send + 'static;
19
20 #[allow(clippy::type_complexity)]
22 fn accept(&mut self) -> Pin<Box<dyn Future<Output = crate::error::Result<IncomingConnection<Self::IO>>> + Send + '_>>;
23
24 fn listen_addrs(&self) -> Vec<SocketAddr> {
26 vec![]
27 }
28}
29
30use tokio::net::TcpListener;
32
33pub struct TcpCaptureSource {
34 listener: TcpListener,
35}
36
37impl TcpCaptureSource {
38 pub fn new(listener: TcpListener) -> Self {
39 Self { listener }
40 }
41}
42
43impl CaptureSource for TcpCaptureSource {
44 type IO = tokio::net::TcpStream;
45
46 fn accept(&mut self) -> Pin<Box<dyn Future<Output = crate::error::Result<IncomingConnection<Self::IO>>> + Send + '_>> {
47 Box::pin(async move {
48 let (stream, client_addr) = self.listener.accept().await?;
49 Ok(IncomingConnection {
50 stream,
51 client_addr,
52 target_addr: None, })
54 })
55 }
56
57 fn listen_addrs(&self) -> Vec<SocketAddr> {
58 if let Ok(addr) = self.listener.local_addr() {
59 vec![addr]
60 } else {
61 vec![]
62 }
63 }
64}