conmonrs/streaming_server/port_forward/
channels.rs

1//! The channel protocol: a connection per port in the query, each with
2//! a data and an error channel.
3
4use super::{
5    CHANNEL_BUFFER, CHANNEL_QUEUE, Early, MAX_CHANNEL_PORTS, OUTPUT_QUEUE, PORT_QUERY, PortForward,
6    namespace::Namespace,
7};
8use crate::streaming_server::{
9    budget::{Budget, Held},
10    channel,
11    idle::Activity,
12    web::{self, ClientMessage, Message, WebSocket},
13};
14use anyhow::{Context, Result, ensure};
15use bytes::{Bytes, BytesMut};
16use futures::{SinkExt, StreamExt};
17use std::{pin::pin, sync::Arc};
18use tokio::{
19    net::TcpStream,
20    select,
21    sync::mpsc,
22    task::{self, JoinSet},
23};
24use tokio_util::sync::CancellationToken;
25use tracing::{debug, trace};
26
27/// Sends the messages of a port to the client.
28#[derive(Clone)]
29pub(super) struct ChannelOutput {
30    pub(super) messages: mpsc::Sender<Message>,
31    pub(super) base64: bool,
32    pub(super) data: u8,
33    pub(super) error: u8,
34}
35
36impl ChannelOutput {
37    /// Send data on a channel, fails once the client went away.
38    async fn send(&self, channel: u8, data: &[u8]) -> Result<()> {
39        let message = channel::message(self.base64, channel, data);
40        self.messages
41            .send(message)
42            .await
43            .map_err(|_| anyhow::anyhow!("WebSocket closed"))
44    }
45}
46
47impl PortForward {
48    /// The ports a channel protocol request forwards, by its query. The
49    /// `port` parameter may repeat and hold several ports separated by
50    /// commas.
51    pub(crate) fn query_ports(query: Option<&str>) -> Result<Vec<u16>> {
52        let mut ports = vec![];
53        for pair in query.unwrap_or_default().split('&') {
54            let (name, value) = pair.split_once('=').unwrap_or((pair, ""));
55            if Self::percent_decode(name) != PORT_QUERY {
56                continue;
57            }
58            let value = Self::percent_decode(value);
59            ensure!(
60                !value.is_empty(),
61                "query parameter {PORT_QUERY:?} cannot be empty"
62            );
63            for port in value.split(',') {
64                let port: u16 = port
65                    .parse()
66                    .with_context(|| format!("unable to parse {port:?} as a port"))?;
67                ensure!(port > 0, "port {port} must be > 0");
68                ports.push(port);
69            }
70        }
71        ensure!(
72            !ports.is_empty(),
73            "query parameter {PORT_QUERY:?} is required"
74        );
75        ensure!(
76            ports.len() <= MAX_CHANNEL_PORTS,
77            "at most {MAX_CHANNEL_PORTS} ports can be forwarded at once"
78        );
79        Ok(ports)
80    }
81
82    /// Decode a query component: a `+` is a space and `%XX` the byte of the
83    /// two hex digits. Invalid escapes stay as they are, and bytes which are
84    /// no UTF-8 become the replacement character.
85    pub(super) fn percent_decode(s: &str) -> String {
86        let hex = |digit: u8| match digit {
87            b'0'..=b'9' => Some(digit - b'0'),
88            b'a'..=b'f' => Some(digit - b'a' + 10),
89            b'A'..=b'F' => Some(digit - b'A' + 10),
90            _ => None,
91        };
92        let mut decoded = Vec::with_capacity(s.len());
93        let mut rest = s.as_bytes();
94        while let [byte, tail @ ..] = rest {
95            rest = tail;
96            decoded.push(match (byte, tail) {
97                (b'%', [high, low, tail @ ..]) => match (hex(*high), hex(*low)) {
98                    (Some(high), Some(low)) => {
99                        rest = tail;
100                        (high << 4) | low
101                    }
102                    _ => b'%',
103                },
104                (b'+', _) => b' ',
105                (byte, _) => *byte,
106            });
107        }
108        String::from_utf8_lossy(&decoded).into_owned()
109    }
110
111    /// Serve the channel protocol: forward a connection per port until all
112    /// of them ended, the client went away or the WebSocket was idle.
113    pub(super) async fn serve_channels(self, socket: WebSocket, ports: Vec<u16>, base64: bool) {
114        let discard = socket.get_ref().discarder();
115        let written = socket.get_ref().written();
116        let (mut sink, mut messages) = socket.split();
117        let (output_tx, mut output_rx) = mpsc::channel(OUTPUT_QUEUE);
118        let activity = Activity::new();
119        let too_large = CancellationToken::new();
120
121        // On its own, as a port which does not take its input stops reading
122        // the WebSocket.
123        let idle = CancellationToken::new();
124        // Ends the watch however the session ends.
125        let _idle_guard = idle.clone().drop_guard();
126        activity.watch(self.idle_timeout, idle.clone(), {
127            let idle = idle.clone();
128            move || {
129                debug!("Closing idle port forward WebSocket");
130                idle.cancel();
131            }
132        });
133
134        let (writer_activity, writer_too_large) = (activity.clone(), too_large.clone());
135        // The sink, once the close went out.
136        let writer = task::spawn(async move {
137            while let Some(message) = output_rx.recv().await {
138                if sink.send(message).await.is_err() {
139                    return None;
140                }
141                writer_activity.touch();
142            }
143            web::send_close(sink, writer_too_large.is_cancelled()).await
144        });
145
146        let namespace = Namespace::new(self.netns_path);
147        let budget = Budget::new(CHANNEL_BUFFER);
148        let mut inputs = Vec::with_capacity(ports.len());
149        let mut forwards = JoinSet::new();
150        for (i, port) in ports.into_iter().enumerate() {
151            #[expect(
152                clippy::cast_possible_truncation,
153                reason = "at most MAX_CHANNEL_PORTS ports"
154            )]
155            let output = ChannelOutput {
156                messages: output_tx.clone(),
157                base64,
158                data: (i * 2) as u8,
159                error: (i * 2 + 1) as u8,
160            };
161            // The client learns the port of each channel from its first
162            // message.
163            let prefix = port.to_le_bytes();
164            if output.send(output.data, &prefix).await.is_err()
165                || output.send(output.error, &prefix).await.is_err()
166            {
167                return;
168            }
169            let (input_tx, input_rx) = mpsc::channel(CHANNEL_QUEUE);
170            inputs.push(input_tx);
171            forwards.spawn(Self::forward_channel(
172                namespace.clone(),
173                port,
174                input_rx,
175                output,
176            ));
177        }
178        drop(output_tx);
179
180        let mut end = web::End::Done;
181        loop {
182            select! {
183                // Nothing gets dispatched after the idle timeout, and the
184                // session ends with its last forward, also while the client
185                // keeps sending.
186                biased;
187                () = idle.cancelled() => break,
188                done = forwards.join_next() => if done.is_none() {
189                    break;
190                },
191                message = web::next_message(&mut messages, &discard) => {
192                    // The heartbeat of a client does not keep it from going
193                    // idle, see [`web::next_message`].
194                    let data = match message {
195                        ClientMessage::Data(message) => message.into_data(),
196                        // The client cannot be read anymore, it gets the
197                        // close only. Like for Kubernetes, that ends the
198                        // session, the error channels of the ports are not
199                        // about it.
200                        ClientMessage::TooLarge => {
201                            too_large.cancel();
202                            end = web::End::TooLarge;
203                            break;
204                        }
205                        ClientMessage::Closed(_) => {
206                            web::answer_close(&mut messages, &written).await;
207                            end = web::End::Gone;
208                            break;
209                        }
210                        ClientMessage::Gone => {
211                            end = web::End::Gone;
212                            break;
213                        }
214                    };
215                    activity.touch();
216                    Self::dispatch(&inputs, &budget, base64, data, &idle).await;
217                    // A dispatch which ended with the idle timeout is no
218                    // activity.
219                    if idle.is_cancelled() {
220                        break;
221                    }
222                    // A port which takes its input slowly keeps the WebSocket
223                    // active.
224                    activity.touch();
225                }
226            }
227        }
228
229        // The ports still forwarded lose their client.
230        forwards.abort_all();
231        while forwards.join_next().await.is_some() {}
232        // Like the error of a port which failed before its input came, the
233        // client has to get what the server sent before.
234        web::end(writer, &written, messages, end).await;
235    }
236
237    /// Hand a message of the client to the connection of its channel. Only
238    /// the data channels, the even ones, take input. Waits while a connected
239    /// port does not take its input or the ports hold all of the budget,
240    /// which stops reading the WebSocket like the reference implementation
241    /// does, the protocol has no flow control. Connecting ports take it, see
242    /// [`Self::drain_until`], and the idle timeout still ends the wait.
243    pub(super) async fn dispatch(
244        inputs: &[mpsc::Sender<Held>],
245        budget: &Budget,
246        base64: bool,
247        message: Bytes,
248        idle: &CancellationToken,
249    ) {
250        let Some((channel, data)) = channel::decode(base64, message) else {
251            return;
252        };
253        let Some(input) = inputs
254            .get(usize::from(channel / 2))
255            .filter(|_| channel % 2 == 0)
256        else {
257            trace!("Dropping port forward message for channel {channel}");
258            return;
259        };
260        if !data.is_empty() {
261            select! {
262                // A connection which ended does not take input anymore.
263                _ = async { input.send(budget.hold(data).await).await } => {}
264                // Also while waiting for the budget.
265                () = input.closed() => {}
266                () = idle.cancelled() => {}
267            }
268        }
269    }
270
271    /// Forward the connection of a port, the error channel reports a failure.
272    pub(super) async fn forward_channel(
273        namespace: Arc<Namespace>,
274        port: u16,
275        mut input: mpsc::Receiver<Held>,
276        output: ChannelOutput,
277    ) {
278        debug!("Forwarding port {port} over WebSocket");
279        let result = match Self::establish(&namespace, port, &mut input).await {
280            Ok((connection, early)) => {
281                Self::copy_channel(early, &mut input, &output, connection).await
282            }
283            Err(e) => Err(e),
284        };
285        // What the client sent gives its budget back, and the WebSocket gets
286        // read on while the error waits for a client which does not read.
287        drop(input);
288        if let Err(e) = result {
289            let message = Self::forward_error(port, &e);
290            debug!("Port forward over WebSocket: {message}");
291            let _ = output.send(output.error, message.as_bytes()).await;
292        }
293    }
294
295    /// Copy between the channels of a port and its connection in the pod,
296    /// until the pod closed the connection, writing to it failed, see
297    /// [`Self::upstream_failed`], or the client went away. The
298    /// protocol has no way to close a channel, the input after the pod closed
299    /// the connection does not matter.
300    async fn copy_channel(
301        early: Early,
302        input: &mut mpsc::Receiver<Held>,
303        output: &ChannelOutput,
304        connection: TcpStream,
305    ) -> Result<()> {
306        let (mut from_container, mut to_container) = connection.into_split();
307        let upstream = Self::write_input(early, input, &mut to_container);
308
309        let downstream = async {
310            let mut buf = BytesMut::new();
311            while let Some(data) = Self::read_pod(&mut from_container, &mut buf).await? {
312                if output.send(output.data, &data).await.is_err() {
313                    // The client went away.
314                    return Ok(());
315                }
316            }
317            anyhow::Ok(())
318        };
319
320        let mut downstream = pin!(downstream);
321        select! {
322            result = &mut downstream => result,
323            result = upstream => match result {
324                Ok(()) => Ok(()),
325                Err(e) => Self::upstream_failed(e, downstream).await,
326            },
327        }
328    }
329}
conmonrs/streaming_server/port_forward/channels.rs

conmonrs/streaming_server/port_forward/
channels.rs
