conmonrs/streaming_server/
attach.rs

1//! Attach sessions, which connect their client to the container process.
2
3use super::{
4    SessionCommon, StreamingServer,
5    session::{ClientRx, InputEvent, SessionInput},
6};
7use crate::{
8    attach::SharedContainerAttach,
9    child_reaper::RegisteredContainer,
10    container_io::{Drain, Lag, Message as IOMessage, Pipe},
11    streaming_server::remotecommand::{MESSAGE_TOO_LARGE, Output, Status},
12};
13use anyhow::{Result, bail};
14use tokio::sync::broadcast::{self, error::RecvError};
15use tokio_util::sync::CancellationToken;
16use tracing::debug;
17
18/// Required attach session data.
19#[derive(Debug)]
20pub(crate) struct AttachSession {
21    pub(crate) child: RegisteredContainer,
22    pub(crate) stop_after_stdin_eof: bool,
23    pub(crate) common: SessionCommon,
24}
25
26impl StreamingServer {
27    /// The attach specific read/write loop. `status` gets the success once
28    /// the session ended, unless the client went away, which goes out after
29    /// the loop.
30    pub(super) async fn attach_loop(
31        session: AttachSession,
32        sender: &mut Output,
33        client: ClientRx,
34        status: &mut Option<Status>,
35    ) -> Result<()> {
36        let AttachSession {
37            child,
38            stop_after_stdin_eof,
39            common:
40                SessionCommon {
41                    stdin: stdin_enabled,
42                    stdout: stdout_enabled,
43                    stderr: stderr_enabled,
44                    ..
45                },
46        } = session;
47        let io = child.io().clone();
48        let token = child.token().clone();
49
50        // Subscribe to the attach broadcast rather than the stdio channels. The
51        // broadcast only carries output produced while someone is attached,
52        // whereas the stdio channels replay buffered history and are closed
53        // for created containers anyway.
54        let attach = io.attach().clone();
55        let terminal = attach.is_terminal();
56        // Without stdin, there is none to wait for.
57        let mut input = SessionInput::new(io, attach.clone(), client, stdin_enabled);
58
59        // With stdin once, the end of the session closes the container stdin,
60        // also if the client went away without closing it, like the first
61        // session ending does with StdinOnce in Kubernetes.
62        let mut teardown = AttachTeardown {
63            attach: attach.clone(),
64            token: token.clone(),
65            close_stdin: stdin_enabled && stop_after_stdin_eof,
66        };
67
68        // A single `Done` follows once all pipes ended. It may have been sent
69        // before subscribing, then only forward what is still queued.
70        let (mut attach_rx, mut done) = attach.subscribe();
71        if done {
72            while let Ok(message) = attach_rx.try_recv() {
73                Self::forward_output(message, stdout_enabled, stderr_enabled, sender).await?;
74            }
75        }
76
77        // Until nothing is left to forward in either direction.
78        let mut lag = Lag::default();
79        while !done || input.stdin_open {
80            tokio::select! {
81                res = attach_rx.recv(), if !done => match res {
82                    Ok(IOMessage::Done) => done = true,
83                    Ok(message) => {
84                        lag.received(&attach_rx);
85                        Self::forward_output(message, stdout_enabled, stderr_enabled, sender).await?;
86                    }
87                    Err(RecvError::Lagged(n)) => lag.skipped(n),
88                    Err(RecvError::Closed) => {
89                        debug!("Exiting streaming attach because container I/O is gone");
90                        break
91                    }
92                },
93
94                event = input.next() => match event {
95                    InputEvent::Resize(data) => input.resize(&data),
96                    InputEvent::Ping => sender.ping().await?,
97                    // Without a terminal and without stdin once, that ends
98                    // the session, like for CRI-O with conmon. Otherwise
99                    // keep forwarding the output until it ends, and with
100                    // stdin once close the container stdin.
101                    InputEvent::StdinClosed if !terminal && !stop_after_stdin_eof => break,
102                    InputEvent::StdinClosed => {
103                        if stop_after_stdin_eof {
104                            teardown.close_stdin = false;
105                            input.close_container_stdin();
106                        }
107                    }
108                    // The client closed the connection, everything it sent
109                    // reached the container. Nobody gets a status anymore.
110                    InputEvent::Ended => return Ok(()),
111                    InputEvent::TooLarge => bail!(MESSAGE_TOO_LARGE),
112                },
113
114                () = token.cancelled() => {
115                    debug!("Exiting streaming attach because token cancelled");
116                    if !done {
117                        Self::drain_attach(&attach, &mut attach_rx, &mut lag, stdout_enabled, stderr_enabled, sender).await?;
118                    }
119                    break
120                }
121            }
122        }
123
124        *status = Some(Status::Success);
125        Ok(())
126    }
127
128    /// Forward a single output message to the client.
129    async fn forward_output(
130        message: IOMessage,
131        stdout_enabled: bool,
132        stderr_enabled: bool,
133        sender: &mut Output,
134    ) -> Result<()> {
135        match message {
136            IOMessage::Data(data, pipe @ Pipe::StdOut) if stdout_enabled => {
137                sender.data(pipe, &data).await
138            }
139            IOMessage::Data(data, pipe @ Pipe::StdErr) if stderr_enabled => {
140                sender.data(pipe, &data).await
141            }
142            IOMessage::Data(..) | IOMessage::Done => Ok(()),
143        }
144    }
145
146    /// Forward the output which is still queued or arrives shortly after the
147    /// container exited, until the output ended or the [`Drain`] gives up,
148    /// see [`Drain::recv_attach`].
149    async fn drain_attach(
150        attach: &SharedContainerAttach,
151        attach_rx: &mut broadcast::Receiver<IOMessage>,
152        lag: &mut Lag,
153        stdout_enabled: bool,
154        stderr_enabled: bool,
155        sender: &mut Output,
156    ) -> Result<()> {
157        let mut drain = Drain::default();
158        while let Some(message) = drain.recv_attach(attach, attach_rx, Some(lag)).await {
159            Self::forward_output(message, stdout_enabled, stderr_enabled, sender).await?;
160        }
161        Ok(())
162    }
163}
164
165/// Closes the container stdin when an attach session with stdin once ends
166/// while the container runs, also if the client went away without closing
167/// its stdin and the loop got aborted. A terminal stays open, see
168/// [`SharedContainerAttach::close_container_stdin`].
169struct AttachTeardown {
170    attach: SharedContainerAttach,
171
172    /// Cancelled once the container exited.
173    token: CancellationToken,
174
175    /// Whether the container stdin still has to be closed.
176    close_stdin: bool,
177}
178
179impl Drop for AttachTeardown {
180    fn drop(&mut self) {
181        if self.close_stdin && !self.token.is_cancelled() && self.attach.close_container_stdin() {
182            debug!("Closed the container stdin after the end of the attach session");
183        }
184    }
185}
186
187#[cfg(test)]
188mod tests;
conmonrs/streaming_server/attach.rs

conmonrs/streaming_server/
attach.rs
