kube_client/api/
remote_command.rs

1use std::time::Duration;
2
3use bytes::Bytes;
4use k8s_openapi::apimachinery::pkg::apis::meta::v1::Status;
5
6use futures::{
7    FutureExt, SinkExt, StreamExt,
8    channel::{mpsc, oneshot},
9};
10use serde::{Deserialize, Serialize};
11use thiserror::Error;
12use tokio::{
13    io::{AsyncRead, AsyncWrite, AsyncWriteExt, DuplexStream},
14    select, time,
15};
16use tokio_tungstenite::tungstenite as ws;
17
18use crate::client::Connection;
19
20use super::AttachParams;
21
22type StatusReceiver = oneshot::Receiver<Status>;
23type StatusSender = oneshot::Sender<Status>;
24
25type TerminalSizeReceiver = mpsc::Receiver<TerminalSize>;
26type TerminalSizeSender = mpsc::Sender<TerminalSize>;
27
28/// TerminalSize define the size of a terminal
29#[derive(Debug, Serialize, Deserialize)]
30#[cfg_attr(docsrs, doc(cfg(feature = "ws")))]
31#[serde(rename_all = "PascalCase")]
32pub struct TerminalSize {
33    /// width of the terminal
34    pub width: u16,
35    /// height of the terminal
36    pub height: u16,
37}
38
39/// Errors from attaching to a pod.
40#[derive(Debug, Error)]
41pub enum Error {
42    /// Failed to read from stdin
43    #[error("failed to read from stdin: {0}")]
44    ReadStdin(#[source] std::io::Error),
45
46    /// Failed to send stdin data to the pod
47    #[error("failed to send a stdin data: {0}")]
48    SendStdin(#[source] ws::Error),
49
50    /// Failed to write to stdout
51    #[error("failed to write to stdout: {0}")]
52    WriteStdout(#[source] std::io::Error),
53
54    /// Failed to write to stderr
55    #[error("failed to write to stderr: {0}")]
56    WriteStderr(#[source] std::io::Error),
57
58    /// Failed to receive a WebSocket message from the server.
59    #[error("failed to receive a WebSocket message: {0}")]
60    ReceiveWebSocketMessage(#[source] ws::Error),
61
62    // Failed to complete the background task
63    #[error("failed to complete the background task: {0}")]
64    Spawn(#[source] tokio::task::JoinError),
65
66    /// Failed to send close message.
67    #[error("failed to send a WebSocket close message: {0}")]
68    SendClose(#[source] ws::Error),
69
70    /// Failed to send ping message.
71    #[error("failed to send a WebSocket ping message: {0}")]
72    SendPing(#[source] ws::Error),
73
74    /// Failed to deserialize status object
75    #[error("failed to deserialize status object: {0}")]
76    DeserializeStatus(#[source] serde_json::Error),
77
78    /// Failed to send status object
79    #[error("failed to send status object")]
80    SendStatus,
81
82    /// Fail to serialize Terminalsize object
83    #[error("failed to serialize TerminalSize object: {0}")]
84    SerializeTerminalSize(#[source] serde_json::Error),
85
86    /// Fail to send terminal size message
87    #[error("failed to send terminal size message: {0}")]
88    SendTerminalSize(#[source] ws::Error),
89
90    /// Failed to set terminal size, tty need to be true to resize the terminal
91    #[error("failed to set terminal size, tty need to be true to resize the terminal")]
92    TtyNeedToBeTrue,
93}
94
95const MAX_BUF_SIZE: usize = 1024;
96
97/// Represents an attached process in a container for [`attach`] and [`exec`].
98///
99/// Provides access to `stdin`, `stdout`, and `stderr` if attached.
100///
101/// Use [`AttachedProcess::join`] to wait for the process to terminate.
102///
103/// [`attach`]: crate::Api::attach
104/// [`exec`]: crate::Api::exec
105#[cfg_attr(docsrs, doc(cfg(feature = "ws")))]
106pub struct AttachedProcess {
107    has_stdin: bool,
108    has_stdout: bool,
109    has_stderr: bool,
110    stdin_writer: Option<DuplexStream>,
111    stdout_reader: Option<DuplexStream>,
112    stderr_reader: Option<DuplexStream>,
113    status_rx: Option<StatusReceiver>,
114    terminal_resize_tx: Option<TerminalSizeSender>,
115    task: Option<tokio::task::JoinHandle<Result<(), Error>>>,
116}
117
118impl Drop for AttachedProcess {
119    fn drop(&mut self) {
120        if let Some(task) = &self.task {
121            task.abort();
122        }
123    }
124}
125
126impl AttachedProcess {
127    pub(crate) fn new(connection: Connection, ap: &AttachParams) -> Self {
128        // To simplify the implementation, always create a pipe for stdin.
129        // The caller does not have access to it unless they had requested.
130        let (stdin_writer, stdin_reader) = tokio::io::duplex(ap.max_stdin_buf_size.unwrap_or(MAX_BUF_SIZE));
131        let (stdout_writer, stdout_reader) = if ap.stdout {
132            let (w, r) = tokio::io::duplex(ap.max_stdout_buf_size.unwrap_or(MAX_BUF_SIZE));
133            (Some(w), Some(r))
134        } else {
135            (None, None)
136        };
137        let (stderr_writer, stderr_reader) = if ap.stderr {
138            let (w, r) = tokio::io::duplex(ap.max_stderr_buf_size.unwrap_or(MAX_BUF_SIZE));
139            (Some(w), Some(r))
140        } else {
141            (None, None)
142        };
143        let (status_tx, status_rx) = oneshot::channel();
144        let (terminal_resize_tx, terminal_resize_rx) = if ap.tty {
145            let (w, r) = mpsc::channel(10);
146            (Some(w), Some(r))
147        } else {
148            (None, None)
149        };
150
151        let task = tokio::spawn(start_message_loop(
152            connection,
153            stdin_reader,
154            stdout_writer,
155            stderr_writer,
156            status_tx,
157            terminal_resize_rx,
158        ));
159
160        AttachedProcess {
161            has_stdin: ap.stdin,
162            has_stdout: ap.stdout,
163            has_stderr: ap.stderr,
164            task: Some(task),
165            stdin_writer: Some(stdin_writer),
166            stdout_reader,
167            stderr_reader,
168            terminal_resize_tx,
169            status_rx: Some(status_rx),
170        }
171    }
172
173    /// Async writer to stdin.
174    /// ```no_run
175    /// # use kube_client::api::AttachedProcess;
176    /// # use tokio::io::{AsyncReadExt, AsyncWriteExt};
177    /// # async fn wrapper() -> Result<(), Box<dyn std::error::Error>> {
178    /// # let attached: AttachedProcess = todo!();
179    /// let mut stdin_writer = attached.stdin().unwrap();
180    /// stdin_writer.write(b"foo\n").await?;
181    /// # Ok(())
182    /// # }
183    /// ```
184    /// Only available if [`AttachParams`](super::AttachParams) had `stdin`.
185    pub fn stdin(&mut self) -> Option<impl AsyncWrite + Unpin + use<>> {
186        if !self.has_stdin {
187            return None;
188        }
189        self.stdin_writer.take()
190    }
191
192    /// Async reader for stdout outputs.
193    /// ```no_run
194    /// # use kube_client::api::AttachedProcess;
195    /// # use tokio::io::{AsyncReadExt, AsyncWriteExt};
196    /// # async fn wrapper() -> Result<(), Box<dyn std::error::Error>> {
197    /// # let attached: AttachedProcess = todo!();
198    /// let mut stdout_reader = attached.stdout().unwrap();
199    /// let mut buf = [0u8; 4];
200    /// stdout_reader.read_exact(&mut buf).await?;
201    /// # Ok(())
202    /// # }
203    /// ```
204    /// Only available if [`AttachParams`](super::AttachParams) had `stdout`.
205    pub fn stdout(&mut self) -> Option<impl AsyncRead + Unpin + use<>> {
206        if !self.has_stdout {
207            return None;
208        }
209        self.stdout_reader.take()
210    }
211
212    /// Async reader for stderr outputs.
213    /// ```no_run
214    /// # use kube_client::api::AttachedProcess;
215    /// # use tokio::io::{AsyncReadExt, AsyncWriteExt};
216    /// # async fn wrapper() -> Result<(), Box<dyn std::error::Error>> {
217    /// # let attached: AttachedProcess = todo!();
218    /// let mut stderr_reader = attached.stderr().unwrap();
219    /// let mut buf = [0u8; 4];
220    /// stderr_reader.read_exact(&mut buf).await?;
221    /// # Ok(())
222    /// # }
223    /// ```
224    /// Only available if [`AttachParams`](super::AttachParams) had `stderr`.
225    pub fn stderr(&mut self) -> Option<impl AsyncRead + Unpin + use<>> {
226        if !self.has_stderr {
227            return None;
228        }
229        self.stderr_reader.take()
230    }
231
232    /// Abort the background task, causing remote command to fail.
233    #[inline]
234    pub fn abort(&self) {
235        if let Some(task) = &self.task {
236            task.abort();
237        }
238    }
239
240    /// Waits for the remote command task to complete.
241    pub async fn join(mut self) -> Result<(), Error> {
242        // Drop all streams before awaiting the task to prevent deadlocks.
243        // If stdout/stderr readers have not been drained, the background task
244        // blocks writing to the full DuplexStream buffer while join() blocks
245        // waiting for the task.
246        self.stdin_writer = None;
247        self.stdout_reader = None;
248        self.stderr_reader = None;
249        self.status_rx = None;
250        self.terminal_resize_tx = None;
251        match self.task.take() {
252            Some(task) => task.await.unwrap_or_else(|e| Err(Error::Spawn(e))),
253            None => Ok(()),
254        }
255    }
256
257    /// Take a future that resolves with any status object or when the sender is dropped.
258    ///
259    /// Returns `None` if called more than once.
260    pub fn take_status(&mut self) -> Option<impl Future<Output = Option<Status>> + use<>> {
261        self.status_rx.take().map(|recv| recv.map(|res| res.ok()))
262    }
263
264    /// Async writer to change the terminal size
265    /// ```no_run
266    /// # use kube_client::api::{AttachedProcess, TerminalSize};
267    /// # use tokio::io::{AsyncReadExt, AsyncWriteExt};
268    /// # use futures::SinkExt;
269    /// # async fn wrapper() -> Result<(), Box<dyn std::error::Error>> {
270    /// # let attached: AttachedProcess = todo!();
271    /// let mut terminal_size_writer = attached.terminal_size().unwrap();
272    /// terminal_size_writer.send(TerminalSize{
273    ///     height: 100,
274    ///     width: 200,
275    /// }).await?;
276    /// # Ok(())
277    /// # }
278    /// ```
279    /// Only available if [`AttachParams`](super::AttachParams) had `tty`.
280    pub fn terminal_size(&mut self) -> Option<TerminalSizeSender> {
281        self.terminal_resize_tx.take()
282    }
283}
284
285// theses values come from here: https://github.com/kubernetes/kubernetes/blob/master/staging/src/k8s.io/apimachinery/pkg/util/remotecommand/constants.go#L57
286const STDIN_CHANNEL: u8 = 0;
287const STDOUT_CHANNEL: u8 = 1;
288const STDERR_CHANNEL: u8 = 2;
289// status channel receives `Status` object on exit.
290const STATUS_CHANNEL: u8 = 3;
291// resize channel is use to send TerminalSize object to change the size of the terminal
292const RESIZE_CHANNEL: u8 = 4;
293/// Used to signal that a channel has reached EOF. Only works on V5 of the protocol.
294const CLOSE_CHANNEL: u8 = 255;
295
296async fn start_message_loop(
297    connection: Connection,
298    stdin: impl AsyncRead + Unpin,
299    mut stdout: Option<impl AsyncWrite + Unpin>,
300    mut stderr: Option<impl AsyncWrite + Unpin>,
301    status_tx: StatusSender,
302    mut terminal_size_rx: Option<TerminalSizeReceiver>,
303) -> Result<(), Error> {
304    let supports_stream_close = connection.supports_stream_close();
305    let stream = connection.into_stream();
306    let mut stdin_stream = tokio_util::io::ReaderStream::new(stdin);
307    let (mut server_send, raw_server_recv) = stream.split();
308    // Work with filtered messages to reduce noise.
309    let mut server_recv = raw_server_recv.filter_map(filter_message).boxed();
310    let mut have_terminal_size_rx = terminal_size_rx.is_some();
311
312    // True until we reach EOF for stdin.
313    let mut stdin_is_open = true;
314
315    let mut ping_interval = time::interval(Duration::from_secs(60));
316    ping_interval.set_missed_tick_behavior(time::MissedTickBehavior::Delay);
317    ping_interval.reset();
318
319    loop {
320        let terminal_size_next = async {
321            match terminal_size_rx.as_mut() {
322                Some(tmp) => Some(tmp.next().await),
323                None => None,
324            }
325        };
326
327        select! {
328            _ = ping_interval.tick() => {
329                // send a ping to keep an idle connection alive
330                server_send
331                    .send(ws::Message::Ping(Bytes::new()))
332                    .await
333                    .map_err(Error::SendPing)?;
334            },
335
336            server_message = server_recv.next() => {
337                match server_message {
338                    Some(Ok(Message::Stdout(bin))) => {
339                        if let Some(stdout) = stdout.as_mut() {
340                            stdout.write_all(&bin[1..]).await.map_err(Error::WriteStdout)?;
341                        }
342                    },
343                    Some(Ok(Message::Stderr(bin))) => {
344                        if let Some(stderr) = stderr.as_mut() {
345                            stderr.write_all(&bin[1..]).await.map_err(Error::WriteStderr)?;
346                        }
347                    },
348                    Some(Ok(Message::Status(bin))) => {
349                        let status = serde_json::from_slice::<Status>(&bin[1..]).map_err(Error::DeserializeStatus)?;
350                        status_tx.send(status).map_err(|_| Error::SendStatus)?;
351                        break
352                    },
353                    Some(Err(err)) => {
354                        return Err(Error::ReceiveWebSocketMessage(err));
355                    },
356                    None => {
357                        // Connection closed properly
358                        break
359                    },
360                }
361            },
362            stdin_message = stdin_stream.next(), if stdin_is_open => {
363                match stdin_message {
364                    Some(Ok(bytes)) => {
365                        if !bytes.is_empty() {
366                            let mut vec = Vec::with_capacity(bytes.len() + 1);
367                            vec.push(STDIN_CHANNEL);
368                            vec.extend_from_slice(&bytes[..]);
369                            server_send
370                                .send(ws::Message::binary(vec))
371                                .await
372                                .map_err(Error::SendStdin)?;
373                        }
374                    },
375                    Some(Err(err)) => {
376                        return Err(Error::ReadStdin(err));
377                    }
378                    None => {
379                        // Stdin closed (writer half dropped).
380                        // Let the server know we reached the end of stdin.
381                        if supports_stream_close {
382                            // Signal stdin has reached EOF.
383                            // See: https://github.com/kubernetes/kubernetes/blob/master/staging/src/k8s.io/apimachinery/pkg/util/httpstream/wsstream/conn.go#L346
384                            let vec = vec![CLOSE_CHANNEL, STDIN_CHANNEL];
385                            server_send
386                                .send(ws::Message::binary(vec))
387                                .await
388                                .map_err(Error::SendStdin)?;
389                        } else {
390                            // Best we can do is trigger the whole websocket to close.
391                            // We may miss out on any remaining stdout data that has not
392                            // been sent yet.
393                            server_send.close().await.map_err(Error::SendClose)?;
394                        }
395
396                        // Do not check stdin_stream for data in future loops.
397                        stdin_is_open = false;
398                    }
399                }
400            },
401            Some(terminal_size_message) = terminal_size_next, if have_terminal_size_rx => {
402                match terminal_size_message {
403                    Some(new_size) => {
404                        let new_size = serde_json::to_vec(&new_size).map_err(Error::SerializeTerminalSize)?;
405                        let mut vec = Vec::with_capacity(new_size.len() + 1);
406                        vec.push(RESIZE_CHANNEL);
407                        vec.extend_from_slice(&new_size[..]);
408                        server_send.send(ws::Message::Binary(vec.into())).await.map_err(Error::SendTerminalSize)?;
409                    },
410                    None => {
411                        have_terminal_size_rx = false;
412                    }
413                }
414            },
415        }
416    }
417
418    Ok(())
419}
420
421/// Channeled messages from the server.
422enum Message {
423    /// To Stdout channel (1)
424    Stdout(Vec<u8>),
425    /// To stderr channel (2)
426    Stderr(Vec<u8>),
427    /// To error/status channel (3)
428    Status(Vec<u8>),
429}
430
431// Filter to reduce all the possible WebSocket messages into a few we expect to receive.
432async fn filter_message(wsm: Result<ws::Message, ws::Error>) -> Option<Result<Message, ws::Error>> {
433    match wsm {
434        // The protocol only sends binary frames.
435        // Message of size 1 (only channel number) is sent on connection.
436        Ok(ws::Message::Binary(bin)) if bin.len() > 1 => match bin[0] {
437            STDOUT_CHANNEL => Some(Ok(Message::Stdout(bin.into()))),
438            STDERR_CHANNEL => Some(Ok(Message::Stderr(bin.into()))),
439            STATUS_CHANNEL => Some(Ok(Message::Status(bin.into()))),
440            // We don't receive messages to stdin and resize channels.
441            _ => None,
442        },
443        // Ignore any other message types.
444        // We can ignore close message because the server never sends anything special.
445        // The connection terminates on `None`.
446        Ok(_) => None,
447        // Fatal errors. `WebSocketStream` turns `ConnectionClosed` and `AlreadyClosed` into `None`.
448        // So these are unrecoverables.
449        Err(err) => Some(Err(err)),
450    }
451}
kube_client/api/remote_command.rs

kube_client/api/
remote_command.rs
