conmonrs/streaming_server/remotecommand/
spdy.rs

1//! Exec and attach over SPDY: the protocol versions, the streams of a
2//! session and the input of the client.
3
4use super::{ClientEvent, ClientInput, Stdin, Version};
5use crate::streaming_server::{
6    spdy::{Connection, ConnectionHandle, Headers, Stream, StreamReceiver, StreamSender},
7    stream_protocol::{ClientStr, STREAM_CREATION_TIMEOUT, STREAM_TYPE_HEADER},
8};
9use anyhow::{Context, Result, bail};
10use bytes::{Bytes, BytesMut};
11use serde::de::IgnoredAny;
12use std::{
13    future::{self, Future},
14    sync::{
15        Arc,
16        atomic::{AtomicBool, Ordering},
17    },
18};
19use tokio::{select, task, time};
20use tokio_util::sync::CancellationToken;
21use tracing::debug;
22
23/// The most resize events a SPDY client may send without completing one.
24pub(super) const MAX_RESIZE_BUFFER: usize = 4 * 1024;
25
26/// The stream type of the resize stream.
27pub(super) const STREAM_TYPE_RESIZE: &str = "resize";
28
29/// The SPDY protocols, the kubelet and client-go prefer the first one.
30pub(crate) const SPDY_PROTOCOLS: [(&str, Version); 4] = [
31    ("v4.channel.k8s.io", Version::V4),
32    ("v3.channel.k8s.io", Version::V3),
33    ("v2.channel.k8s.io", Version::V2),
34    ("channel.k8s.io", Version::V1),
35];
36
37/// The input of a SPDY client: the stdin stream, whose end closes stdin, and
38/// the resize stream, which carries consecutive JSON objects. Clients open a
39/// resize stream only with a terminal, which the server does not know for an
40/// attach, so it may also come after the session started. Both get read
41/// together: once the stdin which the process does not read fills
42/// `PENDING_STDIN_MAX_BYTES` of the session, the resizes wait as well.
43pub(crate) struct SpdyInput {
44    stdin: Option<StreamReceiver>,
45
46    /// Set once the client closed its stdin cleanly, so that the output,
47    /// which fails once the connection ended, does not end the session
48    /// before the stdin got forwarded, see [`Output`](super::Output). The
49    /// end of the connection is a clean close then, which still forwards the
50    /// stdin the client sent before.
51    peer_closed: Arc<AtomicBool>,
52    resize: Option<StreamReceiver>,
53    resize_events: ResizeEvents,
54
55    /// The connection, for the streams which come after the start, see
56    /// [`SpdyInput::late_stream`].
57    connection: Option<Connection>,
58
59    /// Cancelled once the connection ended.
60    pub(super) done: CancellationToken,
61}
62
63/// What a SPDY client sent.
64enum SpdyEvent {
65    Stdin(Option<Bytes>),
66    Resize(Option<Bytes>),
67    Stream(Option<Stream>),
68    Done,
69}
70
71/// The output of `f` for the value, pending without one, like for a stream
72/// which ended.
73async fn or_pending<T, F>(value: Option<T>, f: impl FnOnce(T) -> F) -> F::Output
74where
75    F: Future,
76{
77    match value {
78        Some(value) => f(value).await,
79        None => future::pending().await,
80    }
81}
82
83impl SpdyInput {
84    /// Take a stream which comes after the start: the resize stream, and
85    /// output streams the session does not use, which end right away, like
86    /// stderr which clients open for a terminal without asking for one.
87    /// Others get reset. Nothing gets awaited, the caller may cancel
88    /// [`ClientInput::next`].
89    fn late_stream(&mut self, stream: Stream) {
90        match StreamType::parse(stream.headers()) {
91            Ok(StreamType::Resize) if self.resize.is_none() => {
92                let (receiver, sender) = stream.split();
93                task::spawn(close_sending_side(sender));
94                self.resize = Some(receiver);
95            }
96            Ok(StreamType::Stdout | StreamType::Stderr) => {
97                task::spawn(close_sending_side(stream.split().1));
98            }
99            _ => debug!("Rejecting a stream opened after the start of the session"),
100        }
101    }
102}
103
104/// The events of the resize stream of a SPDY client, which carries JSON
105/// objects back to back. Only the latest one matters, the earlier sizes are
106/// gone already, so a client sending many at once does not queue them.
107#[derive(Debug, Default)]
108pub(super) struct ResizeEvents {
109    pub(super) buf: BytesMut,
110    latest: Option<Bytes>,
111}
112
113impl ResizeEvents {
114    /// Split what the client sent into its JSON objects and keep the last
115    /// one. A broken one drops what the client sent so far.
116    pub(super) fn push(&mut self, data: &[u8]) {
117        self.buf.extend_from_slice(data);
118        let mut stream = serde_json::Deserializer::from_slice(&self.buf).into_iter::<IgnoredAny>();
119        let mut start = 0;
120        let mut latest = None;
121        loop {
122            match stream.next() {
123                Some(Ok(_)) => {
124                    let end = stream.byte_offset();
125                    latest = Some(start..end);
126                    start = end;
127                }
128                Some(Err(e)) if e.is_eof() => break,
129                Some(Err(e)) => {
130                    debug!("Dropping invalid resize events: {e}");
131                    start = self.buf.len();
132                    break;
133                }
134                None => break,
135            }
136        }
137        if let Some(range) = latest {
138            self.latest = Some(Bytes::copy_from_slice(self.buf[range].trim_ascii()));
139        }
140        let _ = self.buf.split_to(start);
141        let whitespace = self
142            .buf
143            .iter()
144            .take_while(|b| b.is_ascii_whitespace())
145            .count();
146        let _ = self.buf.split_to(whitespace);
147        if self.buf.len() > MAX_RESIZE_BUFFER {
148            debug!("Dropping incomplete resize event");
149            self.buf.clear();
150        }
151    }
152
153    /// Take the latest resize event, if one came since the last one got
154    /// taken.
155    pub(super) fn pop(&mut self) -> Option<Bytes> {
156        self.latest.take()
157    }
158}
159
160impl ClientInput for SpdyInput {
161    async fn next(&mut self) -> Option<ClientEvent> {
162        loop {
163            if let Some(event) = self.resize_events.pop() {
164                return Some(ClientEvent::Resize(event));
165            }
166            // The connection is done once its reader ended, which happens
167            // before the data queued on its streams got taken. What the
168            // client sent before closing still goes first. The rare resizes
169            // go before stdin, which a busy client keeps ready.
170            let event = select! {
171                biased;
172                data = or_pending(self.resize.as_mut(), StreamReceiver::recv) => SpdyEvent::Resize(data),
173                data = or_pending(self.stdin.as_mut(), StreamReceiver::recv) => SpdyEvent::Stdin(data),
174                stream = or_pending(self.connection.as_mut(), Connection::accept) => SpdyEvent::Stream(stream),
175                () = self.done.cancelled() => SpdyEvent::Done,
176            };
177            match event {
178                SpdyEvent::Stdin(Some(data)) => {
179                    if let Some(stdin) = Stdin::data(data) {
180                        return Some(ClientEvent::Stdin(stdin));
181                    }
182                }
183                // The client closed stdin, unless the stream got reset, also
184                // if the connection ended right after.
185                SpdyEvent::Stdin(None) => {
186                    let reset = self
187                        .stdin
188                        .take()
189                        .is_some_and(|s| s.is_reset() && !s.is_closed_by_client());
190                    if !reset {
191                        self.peer_closed.store(true, Ordering::Release);
192                        return Some(ClientEvent::Stdin(Stdin::Close));
193                    }
194                }
195                SpdyEvent::Resize(Some(data)) => self.resize_events.push(&data),
196                SpdyEvent::Resize(None) => self.resize = None,
197                SpdyEvent::Stream(Some(stream)) => self.late_stream(stream),
198                SpdyEvent::Stream(None) => self.connection = None,
199                // The session sets it as well, but only once the input ended
200                // with the close, so it tells whether stdin got closed.
201                SpdyEvent::Done => {
202                    return self
203                        .peer_closed
204                        .load(Ordering::Acquire)
205                        .then_some(ClientEvent::Closed);
206                }
207            }
208        }
209    }
210}
211
212/// The receiving side of a stream, whose sending side the server does not
213/// need.
214async fn receiver(stream: Option<Stream>) -> Option<StreamReceiver> {
215    let (receiver, sender) = stream?.split();
216    close_sending_side(sender).await;
217    Some(receiver)
218}
219
220/// Close the sending side of a stream. A late stream gets it in a task of
221/// its own, without waiting for the queue of the connection.
222async fn close_sending_side(mut sender: StreamSender) {
223    if let Err(e) = sender.close().await {
224        debug!("Unable to close the sending side of a stream: {e:#}");
225    }
226}
227
228/// The type of a stream of a session, by its header.
229#[derive(Clone, Copy, Debug, PartialEq)]
230enum StreamType {
231    Error,
232    Stdin,
233    Stdout,
234    Stderr,
235    Resize,
236}
237
238impl StreamType {
239    fn parse(headers: &Headers) -> Result<Self> {
240        match headers.get(STREAM_TYPE_HEADER) {
241            Some("error") => Ok(Self::Error),
242            Some("stdin") => Ok(Self::Stdin),
243            Some("stdout") => Ok(Self::Stdout),
244            Some("stderr") => Ok(Self::Stderr),
245            Some(STREAM_TYPE_RESIZE) => Ok(Self::Resize),
246            Some(typ) => bail!("unexpected stream type {:?}", ClientStr(typ)),
247            None => bail!("{STREAM_TYPE_HEADER:?} header is required"),
248        }
249    }
250}
251
252/// The streams a session expects from a SPDY client. A resize stream is
253/// optional, see [`SpdyInput`].
254#[derive(Clone, Copy, Debug)]
255pub(crate) struct ExpectedStreams {
256    pub(crate) stdin: bool,
257    pub(crate) stdout: bool,
258    pub(crate) stderr: bool,
259}
260
261/// The streams of a SPDY session, which the client opened until the session
262/// started.
263pub(crate) struct SpdyStreams {
264    error: Stream,
265    stdin: Option<Stream>,
266    stdout: Option<Stream>,
267    stderr: Option<Stream>,
268    resize: Option<Stream>,
269}
270
271impl SpdyStreams {
272    /// Whether the client may open a stream, by its headers.
273    pub(crate) fn accept(headers: &Headers, version: Version) -> Result<()> {
274        match StreamType::parse(headers)? {
275            StreamType::Resize if !version.spdy_resize() => {
276                bail!("unexpected stream type {STREAM_TYPE_RESIZE:?}")
277            }
278            _ => Ok(()),
279        }
280    }
281
282    /// Wait for the streams the session expects, others may come along.
283    /// Nothing reads the streams meanwhile: a client which sends more stdin
284    /// than a stream queues before it opened the others holds up the
285    /// connection until this times out. client-go opens all streams first.
286    pub(crate) async fn wait(
287        connection: &mut Connection,
288        expected: ExpectedStreams,
289    ) -> Result<Self> {
290        let mut error = None;
291        let (mut stdin, mut stdout, mut stderr, mut resize) = (None, None, None, None);
292        let timeout = time::sleep(STREAM_CREATION_TIMEOUT);
293        tokio::pin!(timeout);
294        loop {
295            if let Some(error) = error.take_if(|_| {
296                (!expected.stdin || stdin.is_some())
297                    && (!expected.stdout || stdout.is_some())
298                    && (!expected.stderr || stderr.is_some())
299            }) {
300                return Ok(Self {
301                    error,
302                    stdin,
303                    stdout,
304                    stderr,
305                    resize,
306                });
307            }
308            let stream = select! {
309                stream = connection.accept() => {
310                    stream.context("connection ended before the client created all streams")?
311                }
312                () = &mut timeout => bail!("timed out waiting for client to create streams"),
313            };
314            // Only accepted streams come, see [`Self::accept`]. Another one
315            // gets reset like a second stream of a type.
316            let slot = match StreamType::parse(stream.headers()) {
317                Ok(StreamType::Error) => &mut error,
318                Ok(StreamType::Stdin) => &mut stdin,
319                Ok(StreamType::Stdout) => &mut stdout,
320                Ok(StreamType::Stderr) => &mut stderr,
321                Ok(StreamType::Resize) => &mut resize,
322                Err(_) => continue,
323            };
324            if slot.is_some() {
325                // Dropping it resets it.
326                debug!("Rejecting a second stream of the same type");
327                continue;
328            }
329            *slot = Some(stream);
330        }
331    }
332
333    /// The input and the output of the session. The server sends nothing on
334    /// the streams of the client. The sending side of stdin stays open, a
335    /// client takes its end as the end of stdin, except for v1 like for the
336    /// kubelet. The other sending sides get closed.
337    pub(crate) async fn into_parts(
338        self,
339        connection: Connection,
340        version: Version,
341    ) -> (SpdyInput, SpdyOutput) {
342        let (stdin, stdin_sender) = match self.stdin.map(Stream::split) {
343            Some((receiver, sender)) if version == Version::V1 => {
344                close_sending_side(sender).await;
345                (Some(receiver), None)
346            }
347            Some((receiver, sender)) => (Some(receiver), Some(sender)),
348            None => (None, None),
349        };
350        let sender = |stream: Option<Stream>| stream.map(|s| s.split().1);
351        let peer_closed = Arc::<AtomicBool>::default();
352        let output = SpdyOutput {
353            peer_closed: peer_closed.clone(),
354            version,
355            error: self.error.split().1,
356            stdin: stdin_sender,
357            stdout: sender(self.stdout),
358            stderr: sender(self.stderr),
359            connection: connection.handle(),
360        };
361        let input = SpdyInput {
362            peer_closed,
363            stdin,
364            resize: receiver(self.resize).await,
365            resize_events: ResizeEvents::default(),
366            done: connection.handle().done(),
367            connection: Some(connection),
368        };
369        (input, output)
370    }
371}
372
373/// The sending sides of a SPDY session.
374pub(crate) struct SpdyOutput {
375    /// Shared with the input, see [`Output::spdy`](super::Output::spdy).
376    pub(super) peer_closed: Arc<AtomicBool>,
377    pub(super) version: Version,
378    pub(super) error: StreamSender,
379    pub(super) stdin: Option<StreamSender>,
380    pub(super) stdout: Option<StreamSender>,
381    pub(super) stderr: Option<StreamSender>,
382    pub(super) connection: ConnectionHandle,
383}
conmonrs/streaming_server/remotecommand/spdy.rs

conmonrs/streaming_server/remotecommand/
spdy.rs
