nitro_cli/
utils.rs

1// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3#![deny(missing_docs)]
4#![deny(warnings)]
5
6use libc::{c_void, close};
7use nix::poll::poll;
8use nix::poll::{PollFd, PollFlags};
9use nix::sys::socket::{connect, socket};
10use nix::sys::socket::{AddressFamily, SockFlag, SockType, VsockAddr};
11use nix::sys::time::{TimeVal, TimeValLike};
12use nix::unistd::read;
13use std::io::Write;
14use std::mem::size_of;
15use std::os::unix::io::IntoRawFd;
16use std::os::unix::io::RawFd;
17use std::thread::sleep;
18use std::time::{Duration, SystemTime};
19use vmm_sys_util::epoll::{ControlOperation, Epoll, EpollEvent, EventSet};
20use vmm_sys_util::timerfd::TimerFd;
21
22use crate::common::{NitroCliErrorEnum, NitroCliFailure, NitroCliResult};
23use crate::new_nitro_cli_failure;
24
25/// The size of the buffers used for reading console data.
26const BUFFER_SIZE: usize = 1024;
27
28/// The console connection time-out, in milliseconds.
29const CONSOLE_CONNECT_TIMEOUT: i64 = 20000;
30
31/// The `poll` time-out, in milliseconds.
32const POLL_TIMEOUT: i32 = 10000;
33
34/// The socket connection time-out flag.
35const SO_VM_SOCKETS_CONNECT_TIMEOUT: i32 = 6;
36
37/// The amount of time to wait between consecutive console reads, in milliseconds.
38const TIMEOUT: u64 = 100;
39
40/// Defines the types of PCRs that can be measured by `pcr` command
41pub enum PcrType {
42    /// Used for files containing the bytes for hashing
43    DefaultType,
44    /// Used for `.pem` files that we want to hash. Additional serializing is needed
45    SigningCertificate,
46}
47
48/// The structure representing the console of an enclave.
49pub struct Console {
50    /// The file descriptor used for connecting to the enclave's console.
51    fd: RawFd,
52}
53
54impl Drop for Console {
55    fn drop(&mut self) {
56        unsafe { close(self.fd) };
57    }
58}
59
60impl Console {
61    /// Create a new blocking `Console` connection from a given enclave CID and a vsock port.
62    pub fn new(cid: u32, port: u32) -> NitroCliResult<Self> {
63        let socket_fd = socket(
64            AddressFamily::Vsock,
65            SockType::Stream,
66            SockFlag::empty(),
67            None,
68        )
69        .map_err(|err| {
70            new_nitro_cli_failure!(
71                &format!("Failed to create blocking console socket: {err:?}"),
72                NitroCliErrorEnum::SocketError
73            )
74        })?;
75
76        let sockaddr = VsockAddr::new(cid, port);
77
78        vsock_set_connect_timeout(socket_fd, CONSOLE_CONNECT_TIMEOUT).map_err(|err| {
79            err.add_subaction("Failed to set console connect timeout".to_string())
80        })?;
81
82        connect(socket_fd, &sockaddr).map_err(|err| {
83            new_nitro_cli_failure!(
84                &format!("Failed to connect to the console: {err:?}"),
85                NitroCliErrorEnum::EnclaveConsoleConnectionFailure
86            )
87        })?;
88
89        Ok(Console { fd: socket_fd })
90    }
91
92    /// Create a new non-blocking `Console` connection from a given enclave CID and a vsock port.
93    pub fn new_nonblocking(cid: u32, port: u32) -> NitroCliResult<Self> {
94        // create new non blocking socket
95        let socket_fd = socket(
96            AddressFamily::Vsock,
97            SockType::Stream,
98            SockFlag::SOCK_NONBLOCK,
99            None,
100        )
101        .map_err(|err| {
102            new_nitro_cli_failure!(
103                &format!("Failed to create nonblocking console socket: {err:?}"),
104                NitroCliErrorEnum::SocketError
105            )
106        })?;
107
108        vsock_set_connect_timeout(socket_fd, CONSOLE_CONNECT_TIMEOUT).map_err(|err| {
109            err.add_subaction("Failed to set console connect timeout".to_string())
110        })?;
111
112        let sockaddr = VsockAddr::new(cid, port);
113        let result = connect(socket_fd, &sockaddr);
114
115        match result {
116            Ok(_) => println!("Connected to the console"),
117            Err(error) => match error {
118                // If the connection is not ready, wait until socket_fd is ready for writing.
119                nix::errno::Errno::EINPROGRESS => {
120                    let poll_fd = PollFd::new(socket_fd, PollFlags::POLLOUT);
121                    let mut poll_fds = [poll_fd];
122                    match poll(&mut poll_fds, POLL_TIMEOUT) {
123                        Ok(1) => println!("Connected to the console"),
124                        _ => {
125                            return Err(new_nitro_cli_failure!(
126                                "Failed to connect to the console",
127                                NitroCliErrorEnum::SocketError
128                            ))
129                        }
130                    }
131                }
132                _ => {
133                    return Err(new_nitro_cli_failure!(
134                        "Failed to connect to the console",
135                        NitroCliErrorEnum::SocketError
136                    ))
137                }
138            },
139        };
140
141        Ok(Console { fd: socket_fd })
142    }
143
144    /// Read a chunk of raw data from the console and output it.
145    pub fn read_to(
146        &self,
147        output: &mut dyn Write,
148        disconnect_timeout_sec: Option<u64>,
149    ) -> NitroCliResult<()> {
150        // Initialize variables
151        let epoll = Epoll::new().map_err(|e| {
152            new_nitro_cli_failure!(
153                &format!("Failed to create epoll: {e:?}"),
154                NitroCliErrorEnum::EpollError
155            )
156        })?;
157
158        // Add console fd to epoll
159        epoll
160            .ctl(
161                ControlOperation::Add,
162                self.fd,
163                EpollEvent::new(EventSet::IN, self.fd as u64),
164            )
165            .map_err(|e| {
166                new_nitro_cli_failure!(
167                    &format!("Failed to add fd to epoll: {e:?}"),
168                    NitroCliErrorEnum::EpollError
169                )
170            })?;
171
172        // If the function call provides a disconnect timeout, create a timerfd,
173        // arm it and then add it to epoll
174        if let Some(disconnect_timeout) = disconnect_timeout_sec {
175            // Create timerfd
176            let mut timerfd = TimerFd::new().map_err(|e| {
177                new_nitro_cli_failure!(
178                    &format!("Failed to initialize timerfd: {e:?}"),
179                    NitroCliErrorEnum::EpollError
180                )
181            })?;
182
183            // Arm timerfd with disconnect_timeout seconds
184            timerfd
185                .reset(Duration::from_secs(disconnect_timeout), None)
186                .map_err(|e| {
187                    new_nitro_cli_failure!(
188                        &format!("Failed to arm timerfd: {e:?}"),
189                        NitroCliErrorEnum::EpollError
190                    )
191                })?;
192
193            // Add timerfd fd to epoll
194            let timerfd_fd = timerfd.into_raw_fd();
195            epoll
196                .ctl(
197                    ControlOperation::Add,
198                    timerfd_fd,
199                    EpollEvent::new(EventSet::IN, timerfd_fd as u64),
200                )
201                .map_err(|e| {
202                    new_nitro_cli_failure!(
203                        &format!("Failed to add fd to epoll: {e:?}"),
204                        NitroCliErrorEnum::EpollError
205                    )
206                })?;
207        }
208
209        // Allow only one epoll event to happen at a given time
210        let mut events = [EpollEvent::default(); 1];
211
212        loop {
213            // Wait for kernel notification that one of the fds is available
214            let num_events = epoll.wait(-1, &mut events).map_err(|e| {
215                new_nitro_cli_failure!(
216                    &format!("Failed to wait epoll: {e:?}"),
217                    NitroCliErrorEnum::EpollError
218                )
219            })?;
220
221            // Check if any event triggered, because an interrupt could unblock the wait
222            // without any of the requested events to occur
223            if num_events == 1 {
224                match events[0].fd() {
225                    // Check if console fd triggered
226                    fd if fd == self.fd => {
227                        let mut buffer = [0u8; BUFFER_SIZE];
228                        let size = read(self.fd, &mut buffer).map_err(|e| {
229                            new_nitro_cli_failure!(
230                                &format!("Failed to read data from the console: {e:?}"),
231                                NitroCliErrorEnum::EnclaveConsoleReadError
232                            )
233                        })?;
234
235                        if size == 0 {
236                            break;
237                        }
238
239                        if size > 0 {
240                            output.write(&buffer[..size]).map_err(|e| {
241                                new_nitro_cli_failure!(
242                                    &format!(
243                                        "Failed to write data from the \
244                                        console to the given stream: {e:?}"
245                                    ),
246                                    NitroCliErrorEnum::EnclaveConsoleWriteOutputError
247                                )
248                            })?;
249                        }
250                    }
251                    // Check if timerfd triggered
252                    _ => break,
253                }
254            }
255        }
256
257        Ok(())
258    }
259
260    /// Read a chunk of raw data to a buffer.
261    pub fn read_to_buffer(&self, buf: &mut Vec<u8>, duration: Duration) -> NitroCliResult<()> {
262        let sys_time = SystemTime::now();
263
264        loop {
265            let mut buffer = [0u8; BUFFER_SIZE];
266            let result = read(self.fd, &mut buffer);
267
268            if let Ok(size) = result {
269                if size > 0 {
270                    let mut buf_vec = buffer.to_vec();
271                    buf_vec.truncate(size);
272                    (*buf).append(&mut buf_vec);
273                }
274            }
275
276            sleep(Duration::from_millis(TIMEOUT));
277
278            let time_elapsed = sys_time.elapsed().map_err(|err| {
279                new_nitro_cli_failure!(
280                    &format!("System time moved backwards: {err:?}"),
281                    NitroCliErrorEnum::ClockSkewError
282                )
283            })?;
284
285            if time_elapsed >= duration {
286                break;
287            }
288        }
289
290        Ok(())
291    }
292}
293
294/// Set a timeout on a vsock connection.
295fn vsock_set_connect_timeout(fd: RawFd, millis: i64) -> NitroCliResult<()> {
296    let timeval = TimeVal::milliseconds(millis);
297    let ret = unsafe {
298        libc::setsockopt(
299            fd,
300            libc::AF_VSOCK,
301            SO_VM_SOCKETS_CONNECT_TIMEOUT,
302            &timeval as *const _ as *const c_void,
303            size_of::<TimeVal>() as u32,
304        )
305    };
306
307    match ret {
308        0 => Ok(()),
309        _ => Err(new_nitro_cli_failure!(
310            &format!("Failed to configure SO_VM_SOCKETS_CONNECT_TIMEOUT: {ret:?}"),
311            NitroCliErrorEnum::SocketConnectTimeoutError
312        )),
313    }
314}
315
316/// Computes the ceil of `lhs / rhs`. Used for reporting the lower
317/// limit of enclave memory based on the EIF file size.
318pub fn ceil_div(lhs: u64, rhs: u64) -> u64 {
319    if rhs == 0 {
320        return u64::MAX;
321    }
322
323    lhs / rhs
324        + match lhs % rhs {
325            0 => 0,
326            _ => 1,
327        }
328}
nitro_cli/utils.rs

nitro_cli/
utils.rs
