conmonrs/streaming_server/port_forward/
requests.rs

1//! The SPDY protocol: the streams of a request pair up by their headers,
2//! then the request forwards one connection.
3
4use super::{
5    CLOSE_GRACE_PERIOD, Early, HALF_CLOSE_GRACE_PERIOD, PORT_HEADER, PortForward,
6    REQUEST_ID_HEADER, STREAM_TYPE_DATA, STREAM_TYPE_ERROR, namespace::Namespace,
7};
8use crate::streaming_server::{
9    spdy::{Connection, Headers, Stream, StreamReceiver, StreamSender},
10    stream_protocol::{ClientStr, STREAM_CREATION_TIMEOUT, STREAM_TYPE_HEADER},
11};
12use anyhow::{Context, Result, bail, ensure};
13use bytes::BytesMut;
14use futures::{StreamExt, stream::FuturesUnordered};
15use std::{
16    collections::HashMap,
17    future::Future,
18    pin::{Pin, pin},
19    sync::Arc,
20};
21use tokio::{io::AsyncWriteExt, net::TcpStream, select, task, time};
22use tokio_util::{
23    sync::{CancellationToken, DropGuard},
24    time::{DelayQueue, delay_queue::Key},
25};
26use tracing::debug;
27
28/// The streams of a request, until both arrived.
29pub(super) struct PendingRequest {
30    data: Option<Waiting>,
31    error: Option<Waiting>,
32
33    /// The deadline of the request in [`Pending::deadlines`].
34    deadline: Key,
35}
36
37/// A stream which waits for the other one of its request. Its reset gets
38/// watched until it stops waiting, see [`Pending::resets`].
39struct Waiting {
40    stream: Stream,
41    _watch: DropGuard,
42}
43
44impl Waiting {
45    fn is_reset(&self) -> bool {
46        self.stream.is_reset()
47    }
48}
49
50/// The watch of a waiting stream, which names its request once the client
51/// reset it, nothing once the stream stopped waiting.
52type ResetWatch = Pin<Box<dyn Future<Output = Option<Arc<str>>> + Send>>;
53
54/// The requests waiting for their second stream, with their deadlines and
55/// the watches of their streams. All are bounded by the streams of a
56/// connection.
57pub(super) struct Pending {
58    pub(super) requests: HashMap<Arc<str>, PendingRequest>,
59    pub(super) deadlines: DelayQueue<Arc<str>>,
60
61    /// The watches of the waiting streams, which let go of a stream once
62    /// the client reset it. Otherwise the streams of a client which opens
63    /// and resets them would count against the most streams of its
64    /// connection until their deadline, as no new stream which could free
65    /// them gets accepted once there is no room left.
66    pub(super) resets: FuturesUnordered<ResetWatch>,
67}
68
69impl Pending {
70    /// No pending requests.
71    pub(super) fn new() -> Self {
72        Self {
73            requests: HashMap::new(),
74            deadlines: DelayQueue::new(),
75            resets: FuturesUnordered::new(),
76        }
77    }
78
79    /// A stream of a request which waits for the other one, and the watch
80    /// for its reset, see [`Self::resets`].
81    fn wait(request_id: &Arc<str>, stream: Stream) -> (Waiting, ResetWatch) {
82        let waiting = CancellationToken::new();
83        let (reset, stopped) = (stream.reset_watch(), waiting.clone());
84        let request_id = request_id.clone();
85        let watch = Box::pin(async move {
86            select! {
87                biased;
88                () = stopped.cancelled() => None,
89                () = reset => Some(request_id),
90            }
91        });
92        let stream = Waiting {
93            stream,
94            _watch: waiting.drop_guard(),
95        };
96        (stream, watch)
97    }
98
99    /// The request, created with its deadline if it is new.
100    pub(super) fn get_or_insert(&mut self, request_id: &Arc<str>) -> &mut PendingRequest {
101        let Self {
102            requests,
103            deadlines,
104            ..
105        } = self;
106        requests
107            .entry(request_id.clone())
108            .or_insert_with(|| PendingRequest {
109                data: None,
110                error: None,
111                deadline: deadlines.insert(request_id.clone(), STREAM_CREATION_TIMEOUT),
112            })
113    }
114
115    /// Remove a request, with its deadline.
116    pub(super) fn remove(&mut self, request_id: &str) -> Option<PendingRequest> {
117        let request = self.requests.remove(request_id)?;
118        self.deadlines.remove(&request.deadline);
119        Some(request)
120    }
121
122    /// Let go of the streams of a request which the client reset, and of
123    /// the request once it has no stream left.
124    pub(super) fn remove_reset(&mut self, request_id: &str) {
125        let Some(request) = self.requests.get_mut(request_id) else {
126            return;
127        };
128        request.data.take_if(|stream| stream.is_reset());
129        request.error.take_if(|stream| stream.is_reset());
130        if request.data.is_none() && request.error.is_none() {
131            self.remove(request_id);
132        }
133    }
134}
135
136impl PortForward {
137    /// Serve port forward requests on the upgraded connection until it ends.
138    pub(crate) async fn serve_spdy<T>(self, io: T)
139    where
140        T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + 'static,
141    {
142        let mut connection =
143            Connection::serve(io, self.idle_timeout, self.ping_interval, Self::accept);
144        let namespace = Namespace::new(self.netns_path);
145        let mut pending = Pending::new();
146
147        loop {
148            select! {
149                Some(reset) = pending.resets.next() => if let Some(request_id) = reset {
150                    pending.remove_reset(&request_id);
151                },
152                stream = connection.accept() => match stream {
153                    Some(stream) => Self::add(&mut pending, stream, &namespace),
154                    None => break,
155                },
156                Some(expired) = pending.deadlines.next() => {
157                    Self::expire(&mut pending, expired.into_inner());
158                }
159            }
160        }
161        debug!("Port forward connection ended");
162    }
163
164    /// Whether a stream is part of a valid request, by its headers.
165    pub(super) fn accept(headers: &Headers) -> Result<()> {
166        Self::port(headers)?;
167        match headers.get(STREAM_TYPE_HEADER) {
168            Some(STREAM_TYPE_DATA | STREAM_TYPE_ERROR) => Ok(()),
169            Some(typ) => bail!("invalid stream type {:?}", ClientStr(typ)),
170            None => bail!("{STREAM_TYPE_HEADER:?} header is required"),
171        }
172    }
173
174    fn port(headers: &Headers) -> Result<u16> {
175        let port = headers
176            .get(PORT_HEADER)
177            .with_context(|| format!("{PORT_HEADER:?} header is required"))?;
178        let port: u16 = port
179            .parse()
180            .with_context(|| format!("unable to parse {:?} as a port", ClientStr(port)))?;
181        ensure!(port > 0, "port {port} must be > 0");
182        Ok(port)
183    }
184
185    fn is_error_stream(stream: &Stream) -> bool {
186        stream.headers().get(STREAM_TYPE_HEADER) == Some(STREAM_TYPE_ERROR)
187    }
188
189    /// The request a stream belongs to. Old clients send no ID, their
190    /// requests use the IDs of the error streams, the data streams follow
191    /// them.
192    fn request_id(stream: &Stream) -> Arc<str> {
193        match stream.headers().get(REQUEST_ID_HEADER) {
194            Some(id) if !id.is_empty() => id.into(),
195            _ if Self::is_error_stream(stream) => stream.id().to_string().into(),
196            // Only a key of the pending requests, like the kubelet does
197            // it. The -1 of stream 1, which has no error stream before it,
198            // matches none, so it waits alone until its deadline.
199            _ => (i64::from(stream.id()) - 2).to_string().into(),
200        }
201    }
202
203    /// Add a stream to its request, which starts once both streams arrived.
204    pub(super) fn add(pending: &mut Pending, stream: Stream, namespace: &Arc<Namespace>) {
205        let request_id = Self::request_id(&stream);
206        let is_error = Self::is_error_stream(&stream);
207        let (stream, watch) = Pending::wait(&request_id, stream);
208        let request = pending.get_or_insert(&request_id);
209        // A stream the client reset does not count, also before the reset
210        // got handled.
211        request.data.take_if(|stream| stream.is_reset());
212        request.error.take_if(|stream| stream.is_reset());
213        let slot = if is_error {
214            &mut request.error
215        } else {
216            &mut request.data
217        };
218        if slot.is_some() {
219            // Dropping it resets it.
220            debug!(
221                "Rejecting stream {} of port forward request {:?}, the request has one of its type already",
222                stream.stream.id(),
223                ClientStr(&request_id)
224            );
225            return;
226        }
227        *slot = Some(stream);
228
229        if request.data.is_none() || request.error.is_none() {
230            pending.resets.push(watch);
231            return;
232        }
233        if let Some(PendingRequest {
234            data: Some(data),
235            error: Some(error),
236            ..
237        }) = pending.remove(&request_id)
238        {
239            task::spawn(Self::forward(
240                request_id,
241                data.stream,
242                error.stream,
243                namespace.clone(),
244            ));
245        }
246    }
247
248    /// Fail a request whose second stream did not arrive in time.
249    fn expire(pending: &mut Pending, request_id: Arc<str>) {
250        let Some(mut request) = pending.requests.remove(&request_id) else {
251            return;
252        };
253        debug!(
254            "Port forward request {:?} timed out waiting for streams",
255            ClientStr(&request_id)
256        );
257        if let Some((_, mut error)) = request.error.take().map(|error| error.stream.split()) {
258            let message = format!("request {request_id}: timed out waiting for streams");
259            task::spawn(async move {
260                let _ = error.send(message.into()).await;
261                let _ = error.close().await;
262            });
263        }
264    }
265
266    /// Forward a single connection, the error stream reports a failure.
267    async fn forward(request_id: Arc<str>, data: Stream, error: Stream, namespace: Arc<Namespace>) {
268        let port = Self::port(data.headers()).unwrap_or_default();
269        let id = ClientStr(&request_id);
270        debug!("Forwarding port {port} for request {id:?}");
271        // The server does not read the error stream, what a client sends on
272        // it gets dropped instead of holding up the connection.
273        let (_, mut error) = error.split();
274
275        let (mut from_client, to_client) = data.split();
276        let result = match Self::establish(&namespace, port, &mut from_client).await {
277            Ok((connection, early)) => Self::copy(early, from_client, to_client, connection).await,
278            Err(e) => {
279                // Dropping the sending side resets the data stream, what the
280                // client sent on it gives its budget back.
281                drop((from_client, to_client));
282                Err(e)
283            }
284        };
285        match result {
286            Ok(()) => debug!("Done forwarding port {port} for request {id:?}"),
287            Err(e) => {
288                let message = Self::forward_error(port, &e);
289                debug!("Port forward request {id:?}: {message}");
290                if let Err(e) = error.send(message.into()).await {
291                    debug!("Unable to report port forward error: {e:#}");
292                }
293            }
294        }
295        if let Err(e) = error.close().await {
296            debug!("Unable to close port forward error stream: {e:#}");
297        }
298    }
299
300    /// Copy between the data stream and the connection in the pod, until
301    /// both sides closed or either reset. A failure resets the data stream,
302    /// and so does a side which does not close within its grace period
303    /// after the other one did.
304    async fn copy(
305        early: Early,
306        mut from_client: StreamReceiver,
307        mut to_client: StreamSender,
308        connection: TcpStream,
309    ) -> Result<()> {
310        let (mut from_container, mut to_container) = connection.into_split();
311
312        let upstream = async {
313            Self::write_input(early, &mut from_client, &mut to_container).await?;
314            if !from_client.is_reset() {
315                to_container
316                    .shutdown()
317                    .await
318                    .context("close the connection to the pod")?;
319            }
320            anyhow::Ok(())
321        };
322
323        let downstream = async {
324            let mut buf = BytesMut::new();
325            loop {
326                let data = select! {
327                    data = Self::read_pod(&mut from_container, &mut buf) => data?,
328                    // Nothing takes the output anymore.
329                    () = to_client.reset_by_peer() => return Ok(()),
330                };
331                let Some(data) = data else {
332                    break;
333                };
334                if let Err(e) = to_client.send(data).await {
335                    return Self::unless_reset(&to_client, e);
336                }
337            }
338            match to_client.close().await {
339                Ok(()) => Ok(()),
340                Err(e) => Self::unless_reset(&to_client, e),
341            }
342        };
343
344        let mut upstream = pin!(upstream);
345        let mut downstream = pin!(downstream);
346        select! {
347            result = &mut downstream => {
348                result?;
349                if time::timeout(CLOSE_GRACE_PERIOD, upstream).await.is_err() {
350                    debug!("Client did not close its side of the connection");
351                }
352                Ok(())
353            }
354            // The client closed its side, the pod may still answer.
355            result = &mut upstream => match result {
356                Ok(()) => {
357                    match time::timeout(HALF_CLOSE_GRACE_PERIOD, downstream).await {
358                        Ok(result) => result,
359                        Err(_) => {
360                            debug!("The pod did not close its side of the connection");
361                            Ok(())
362                        }
363                    }
364                }
365                Err(e) => Self::upstream_failed(e, downstream).await,
366            }
367        }
368    }
369
370    /// A stream the client reset or whose connection ended is no failure.
371    fn unless_reset(sender: &StreamSender, e: anyhow::Error) -> Result<()> {
372        if sender.is_reset() { Ok(()) } else { Err(e) }
373    }
374}
conmonrs/streaming_server/port_forward/requests.rs

conmonrs/streaming_server/port_forward/
requests.rs
