wstunnel/src/tunnel/client/cnx_pool.rs
Σrebe - Romain GERARD f149b8190b
cleanup
2024-07-31 22:12:55 +02:00

74 lines
2 KiB
Rust

use crate::protocols;
use crate::protocols::tls;
use crate::tunnel::client::l4_transport_stream::TransportStream;
use crate::tunnel::client::WsClientConfig;
use async_trait::async_trait;
use bb8::ManageConnection;
use std::ops::Deref;
use std::sync::Arc;
use tracing::instrument;
#[derive(Clone)]
pub struct WsConnection(Arc<WsClientConfig>);
impl WsConnection {
pub fn new(config: Arc<WsClientConfig>) -> Self {
Self(config)
}
}
impl Deref for WsConnection {
type Target = WsClientConfig;
fn deref(&self) -> &Self::Target {
&self.0
}
}
#[async_trait]
impl ManageConnection for WsConnection {
type Connection = Option<TransportStream>;
type Error = anyhow::Error;
#[instrument(level = "trace", name = "cnx_server", skip_all)]
async fn connect(&self) -> Result<Self::Connection, Self::Error> {
let so_mark = self.socket_so_mark;
let timeout = self.timeout_connect;
let tcp_stream = if let Some(http_proxy) = &self.http_proxy {
protocols::tcp::connect_with_http_proxy(
http_proxy,
self.remote_addr.host(),
self.remote_addr.port(),
so_mark,
timeout,
&self.dns_resolver,
)
.await?
} else {
protocols::tcp::connect(
self.remote_addr.host(),
self.remote_addr.port(),
so_mark,
timeout,
&self.dns_resolver,
)
.await?
};
if self.remote_addr.tls().is_some() {
let tls_stream = tls::connect(self, tcp_stream).await?;
Ok(Some(TransportStream::Tls(tls_stream)))
} else {
Ok(Some(TransportStream::Plain(tcp_stream)))
}
}
async fn is_valid(&self, _conn: &mut Self::Connection) -> Result<(), Self::Error> {
Ok(())
}
fn has_broken(&self, conn: &mut Self::Connection) -> bool {
conn.is_none()
}
}