conmonrs/streaming_server/port_forward/
tunnel.rs

1//! Port forwarding over a WebSocket: the protocol, and SPDY tunneled
2//! through binary messages.
3
4use super::{
5    BASE64_CHANNEL_PROTOCOL, CHANNEL_PROTOCOL, PortForward, SPDY_TUNNEL_PROTOCOL, TUNNEL_BUFFER,
6};
7use crate::streaming_server::{
8    stream_protocol::first_supported,
9    web::{self, ClientMessage, Message, WebSocket},
10};
11use bytes::BytesMut;
12use futures::{SinkExt, StreamExt};
13use std::pin::pin;
14use tokio::{
15    io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, DuplexStream},
16    select, task, time,
17};
18use tokio_tungstenite::WebSocketStream;
19use tokio_util::sync::CancellationToken;
20use tracing::debug;
21
22/// A protocol a client may forward ports with over a WebSocket.
23#[derive(Clone, Copy, Debug, PartialEq)]
24pub(crate) enum WebSocketProtocol {
25    /// SPDY tunneled through binary messages.
26    SpdyTunnel,
27
28    /// The channel protocol, its payloads base64 encoded in text messages or
29    /// as they are in binary ones.
30    Channels { base64: bool },
31}
32
33/// The subprotocols of the channel protocol.
34const CHANNEL_PROTOCOLS: [(&str, WebSocketProtocol); 2] = [
35    (
36        CHANNEL_PROTOCOL,
37        WebSocketProtocol::Channels { base64: false },
38    ),
39    (
40        BASE64_CHANNEL_PROTOCOL,
41        WebSocketProtocol::Channels { base64: true },
42    ),
43];
44
45impl WebSocketProtocol {
46    /// The protocol for the subprotocols the client requested, with the one
47    /// the server answers with. The tunnel wins like for the kubelet,
48    /// otherwise the first supported one, none means binary channels.
49    pub(crate) fn negotiate<'a>(
50        requested: impl IntoIterator<Item = &'a str>,
51    ) -> Option<(Self, Option<&'static str>)> {
52        let requested: Vec<_> = requested.into_iter().collect();
53        if requested.contains(&SPDY_TUNNEL_PROTOCOL) {
54            return Some((Self::SpdyTunnel, Some(SPDY_TUNNEL_PROTOCOL)));
55        }
56        if requested.is_empty() {
57            return Some((Self::Channels { base64: false }, None));
58        }
59        first_supported(requested, &CHANNEL_PROTOCOLS).map(|(protocol, p)| (protocol, Some(p)))
60    }
61}
62
63impl PortForward {
64    /// Serve port forwarding on the upgraded WebSocket until it ends. The
65    /// channel protocol forwards the ports, the tunnel ignores them.
66    pub(crate) async fn serve_websocket(
67        self,
68        socket: WebSocket,
69        protocol: WebSocketProtocol,
70        ports: Vec<u16>,
71    ) {
72        match protocol {
73            WebSocketProtocol::SpdyTunnel => {
74                let (server, tunnel) = tokio::io::duplex(TUNNEL_BUFFER);
75                let spdy = task::spawn(self.serve_spdy(server));
76                Self::tunnel(socket, tunnel).await;
77                let _ = spdy.await;
78            }
79            WebSocketProtocol::Channels { base64 } => {
80                self.serve_channels(socket, ports, base64).await;
81            }
82        }
83        debug!("Port forward WebSocket ended");
84    }
85
86    /// Carry the bytes of the SPDY connection in binary messages, until
87    /// either side ends.
88    pub(super) async fn tunnel<T>(
89        socket: WebSocketStream<web::ClientIo<T>>,
90        connection: DuplexStream,
91    ) where
92        T: AsyncRead + AsyncWrite + Unpin + Send + 'static,
93    {
94        let discard = socket.get_ref().discarder();
95        let written = socket.get_ref().written();
96        let (mut sink, mut messages) = socket.split();
97        let (mut from_server, mut to_server) = tokio::io::split(connection);
98        // Cancelled once the client sent a message larger than the server
99        // takes, the close tells it why.
100        let too_large = CancellationToken::new();
101
102        // The sink, once the close went out.
103        let mut to_client = task::spawn({
104            let too_large = too_large.clone();
105            async move {
106                let mut buf = BytesMut::new();
107                loop {
108                    buf.reserve(TUNNEL_BUFFER);
109                    match from_server.read_buf(&mut buf).await {
110                        Ok(0) | Err(_) => break,
111                        Ok(_) => {
112                            if sink
113                                .send(Message::Binary(buf.split().freeze()))
114                                .await
115                                .is_err()
116                            {
117                                return None;
118                            }
119                        }
120                    }
121                }
122                web::send_close(sink, too_large.is_cancelled()).await
123            }
124        });
125
126        let from_client = async {
127            loop {
128                match web::next_message(&mut messages, &discard).await {
129                    ClientMessage::Data(Message::Binary(data)) => {
130                        // The SPDY connection ended, its close goes to the
131                        // client, which still gets read until it answered,
132                        // for a bounded time.
133                        if to_server.write_all(&data).await.is_err() {
134                            web::await_close_answer(&mut messages).await;
135                            break;
136                        }
137                    }
138                    ClientMessage::Data(_) => {
139                        debug!("Ending port forward tunnel, the client sent a text message");
140                        break;
141                    }
142                    ClientMessage::TooLarge => {
143                        too_large.cancel();
144                        break;
145                    }
146                    ClientMessage::Closed(_) => {
147                        web::answer_close(&mut messages, &written).await;
148                        break;
149                    }
150                    ClientMessage::Gone => break,
151                }
152            }
153            // Ends the SPDY connection, which ends the output as well.
154            let _ = to_server.shutdown().await;
155        };
156
157        // The sink once the SPDY connection ended first, none if the client
158        // ended the tunnel.
159        let server_done = {
160            let mut from_client = pin!(from_client);
161            select! {
162                () = &mut from_client => None,
163                sink = &mut to_client => {
164                    let sink = sink.ok().flatten();
165                    // The client answers the close. It still gets read for
166                    // that, a message which is too large ends the wait.
167                    if sink.is_some() {
168                        let _ = time::timeout(web::CLOSE_TIMEOUT, from_client).await;
169                    }
170                    Some(sink)
171                }
172            }
173        };
174        let end = if too_large.is_cancelled() {
175            web::End::TooLarge
176        } else {
177            web::End::Gone
178        };
179        match server_done {
180            None => web::end(to_client, &written, messages, end).await,
181            // A client which sent a message which was too large cannot answer.
182            Some(Some(sink)) if end == web::End::TooLarge => web::drain(sink, messages).await,
183            Some(_) => {}
184        }
185    }
186}
conmonrs/streaming_server/port_forward/tunnel.rs

conmonrs/streaming_server/port_forward/
tunnel.rs
