trz_gateway_server/server/
tunnel.rs

1use std::future::Ready;
2use std::future::ready;
3use std::io::ErrorKind;
4use std::sync::Arc;
5
6use axum::extract::Path;
7use axum::extract::WebSocketUpgrade;
8use axum::extract::ws;
9use axum::http::Uri;
10use axum::response::IntoResponse;
11use bytes::Bytes;
12use hyper_util::rt::TokioIo;
13use nameth::NamedEnumValues as _;
14use nameth::nameth;
15use rustls::pki_types::DnsName;
16use rustls::pki_types::InvalidDnsNameError;
17use rustls::pki_types::ServerName;
18use tokio::io::AsyncRead;
19use tokio::io::AsyncWrite;
20use tonic::transport::Channel;
21use tracing::Instrument as _;
22use tracing::Span;
23use tracing::info;
24use tracing::info_span;
25use tracing::warn;
26use trz_gateway_common::id::ClientId;
27use trz_gateway_common::id::ClientName;
28use trz_gateway_common::is_global::IsGlobalError;
29use trz_gateway_common::to_async_io::WebSocketIo;
30
31use super::Server;
32
33impl Server {
34    /// API to serve tunnel connections
35    ///
36    /// Endpoint: "/remote/tunnel/{client_name}"
37    pub async fn tunnel(
38        self: Arc<Self>,
39        client_id: Option<ClientId>,
40        Path(client_name): Path<ClientName>,
41        web_socket: WebSocketUpgrade,
42    ) -> impl IntoResponse {
43        let span = if let Some(client_id) = client_id {
44            info_span!("Tunnel", %client_name, %client_id)
45        } else {
46            info_span!("Tunnel", %client_name)
47        };
48        let _entered = span.clone().entered();
49        info!("Incoming tunnel");
50        web_socket.on_upgrade(move |web_socket| {
51            let _entered = span.clone().entered();
52            self.process_websocket(client_name, web_socket);
53            ready(())
54        })
55    }
56
57    fn process_websocket(self: Arc<Self>, client_name: ClientName, web_socket: ws::WebSocket) {
58        let (stream, _) = AxumWebSocketIo::to_async_io(web_socket);
59        tokio::spawn(
60            async {
61                self.process_connection(client_name, stream)
62                    .await
63                    .inspect_err(|error| warn!("Failed: {error}"))
64            }
65            .in_current_span(),
66        );
67    }
68
69    async fn process_connection(
70        self: Arc<Self>,
71        client_name: ClientName,
72        connection: impl AsyncRead + AsyncWrite + Unpin + Send + Sync + 'static,
73    ) -> Result<(), TunnelError> {
74        let tls_stream = self
75            .tls_client
76            .get()
77            .map_err(TunnelError::IssuerConfig)?
78            .connect(
79                ServerName::DnsName(DnsName::try_from(client_name.to_string())?),
80                connection,
81            )
82            .await
83            .map_err(TunnelError::TlsConnectError)?;
84
85        // The endpoint is irrelevant: gRPC isn't actually connecting to this endpoint.
86        // Instead we are manually providing the connection using 'connect_with_connector'.
87        // The connection used by gRPC is the bi-directional stream based on the WebSocket.
88        let endpoint = tonic::transport::Endpoint::new(format!(
89            "https://localhost/remote/tunnel/{client_name}"
90        ))
91        .map_err(|_| TunnelError::InvalidEndpoint)?;
92        let connector = make_single_use_connector(tls_stream)
93            .await
94            .map_err(TunnelError::SingleUseConnectorError)?;
95        let channel: Channel = endpoint
96            .connect_with_connector(tower::service_fn(connector))
97            .await
98            .map_err(TunnelError::GrpcConnectError)?;
99
100        self.connections.add(client_name, channel);
101        Ok(())
102    }
103}
104
105async fn make_single_use_connector<S: AsyncRead + AsyncWrite>(
106    stream: S,
107) -> std::io::Result<impl FnMut(Uri) -> Ready<std::io::Result<TokioIo<S>>>> {
108    let span = Span::current();
109    let mut single_use_connection = Some(TokioIo::new(stream));
110    let connector = move |_uri| {
111        span.in_scope(|| {
112            let Some(connection) = single_use_connection.take() else {
113                let error = std::io::Error::new(
114                    ErrorKind::AddrInUse,
115                    "The WebSocket was already used to create a channel",
116                );
117                warn!("{error}");
118                return ready(Err(error));
119            };
120            // `single_use_connection` has been consumed and is now empty.
121            assert!(single_use_connection.is_none());
122            ready(Ok(connection))
123        })
124    };
125    Ok(connector)
126}
127
128#[nameth]
129#[derive(thiserror::Error, Debug)]
130pub enum TunnelError {
131    #[error("[{n}] Failed to create synthetic endpoint", n = self.name())]
132    InvalidEndpoint,
133
134    #[error("[{n}] {0}", n = self.name())]
135    InvalidDnsName(#[from] InvalidDnsNameError),
136
137    #[error("[{n}] {0}", n = self.name())]
138    TlsConnectError(std::io::Error),
139
140    #[error("[{n}] {0}", n = self.name())]
141    SingleUseConnectorError(std::io::Error),
142
143    #[error("[{n}] {0}", n = self.name())]
144    GrpcConnectError(tonic::transport::Error),
145
146    #[error("[{n}] {0}", n = self.name())]
147    IssuerConfig(Arc<dyn IsGlobalError>),
148}
149
150struct AxumWebSocketIo;
151impl WebSocketIo for AxumWebSocketIo {
152    type Message = ws::Message;
153    type Error = axum::Error;
154
155    fn into_data(message: Self::Message) -> Bytes {
156        message.into_data()
157    }
158
159    fn into_messsge(bytes: Bytes) -> Self::Message {
160        ws::Message::Binary(bytes)
161    }
162}
trz_gateway_server/server/tunnel.rs

trz_gateway_server/server/
tunnel.rs
