relay_core_lib/capture/
source.rs1use std::future::Future;
2use std::net::SocketAddr;
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(
23 &mut self,
24 ) -> Pin<Box<dyn Future<Output = crate::error::Result<IncomingConnection<Self::IO>>> + Send + '_>>;
25
26 fn listen_addrs(&self) -> Vec<SocketAddr> {
28 vec![]
29 }
30}
31
32use tokio::net::TcpListener;
34
35pub struct TcpCaptureSource {
36 listener: TcpListener,
37}
38
39impl TcpCaptureSource {
40 pub fn new(listener: TcpListener) -> Self {
41 Self { listener }
42 }
43}
44
45impl CaptureSource for TcpCaptureSource {
46 type IO = tokio::net::TcpStream;
47
48 fn accept(
49 &mut self,
50 ) -> Pin<Box<dyn Future<Output = crate::error::Result<IncomingConnection<Self::IO>>> + Send + '_>>
51 {
52 Box::pin(async move {
53 let (stream, client_addr) = self.listener.accept().await?;
54 Ok(IncomingConnection {
55 stream,
56 client_addr,
57 target_addr: None, })
59 })
60 }
61
62 fn listen_addrs(&self) -> Vec<SocketAddr> {
63 if let Ok(addr) = self.listener.local_addr() {
64 vec![addr]
65 } else {
66 vec![]
67 }
68 }
69}