xref: /cloud-hypervisor/test_infra/src/lib.rs (revision f7f2f25a574b1b2dba22c094fc8226d404157d15)
1 // Copyright © 2021 Intel Corporation
2 //
3 // SPDX-License-Identifier: Apache-2.0
4 //
5 
6 use ssh2::Session;
7 use std::fs;
8 use std::io;
9 use std::io::{Read, Write};
10 use std::net::TcpListener;
11 use std::net::TcpStream;
12 use std::os::unix::io::AsRawFd;
13 use std::path::Path;
14 use std::process::{ExitStatus, Output};
15 use std::str::FromStr;
16 use std::thread;
17 use vmm_sys_util::tempdir::TempDir;
18 
19 pub const DEFAULT_TCP_LISTENER_MESSAGE: &str = "booted";
20 
21 pub struct GuestNetworkConfig {
22     pub guest_ip: String,
23     pub l2_guest_ip1: String,
24     pub l2_guest_ip2: String,
25     pub l2_guest_ip3: String,
26     pub host_ip: String,
27     pub guest_mac: String,
28     pub l2_guest_mac1: String,
29     pub l2_guest_mac2: String,
30     pub l2_guest_mac3: String,
31     pub tcp_listener_port: u16,
32 }
33 
34 pub const DEFAULT_TCP_LISTENER_PORT: u16 = 8000;
35 pub const DEFAULT_TCP_LISTENER_TIMEOUT: i32 = 80;
36 
37 #[derive(Debug)]
38 pub enum WaitForBootError {
39     EpollWait(std::io::Error),
40     Listen(std::io::Error),
41     EpollWaitTimeout,
42     WrongGuestAddr,
43     Accept(std::io::Error),
44 }
45 
46 impl GuestNetworkConfig {
47     pub fn wait_vm_boot(&self, custom_timeout: Option<i32>) -> Result<(), WaitForBootError> {
48         let start = std::time::Instant::now();
49         // The 'port' is unique per 'GUEST' and listening to wild-card ip avoids retrying on 'TcpListener::bind()'
50         let listen_addr = format!("0.0.0.0:{}", self.tcp_listener_port);
51         let expected_guest_addr = self.guest_ip.as_str();
52         let mut s = String::new();
53         let timeout = match custom_timeout {
54             Some(t) => t,
55             None => DEFAULT_TCP_LISTENER_TIMEOUT,
56         };
57 
58         match (|| -> Result<(), WaitForBootError> {
59             let listener =
60                 TcpListener::bind(&listen_addr.as_str()).map_err(WaitForBootError::Listen)?;
61             listener
62                 .set_nonblocking(true)
63                 .expect("Cannot set non-blocking for tcp listener");
64 
65             // Reply on epoll w/ timeout to wait for guest connections faithfully
66             let epoll_fd = epoll::create(true).expect("Cannot create epoll fd");
67             epoll::ctl(
68                 epoll_fd,
69                 epoll::ControlOptions::EPOLL_CTL_ADD,
70                 listener.as_raw_fd(),
71                 epoll::Event::new(epoll::Events::EPOLLIN, 0),
72             )
73             .expect("Cannot add 'tcp_listener' event to epoll");
74             let mut events = vec![epoll::Event::new(epoll::Events::empty(), 0); 1];
75             loop {
76                 let num_events = match epoll::wait(epoll_fd, timeout * 1000_i32, &mut events[..]) {
77                     Ok(num_events) => Ok(num_events),
78                     Err(e) => match e.raw_os_error() {
79                         Some(libc::EAGAIN) | Some(libc::EINTR) => continue,
80                         _ => Err(e),
81                     },
82                 }
83                 .map_err(WaitForBootError::EpollWait)?;
84                 if num_events == 0 {
85                     return Err(WaitForBootError::EpollWaitTimeout);
86                 }
87                 break;
88             }
89 
90             match listener.accept() {
91                 Ok((_, addr)) => {
92                     // Make sure the connection is from the expected 'guest_addr'
93                     if addr.ip() != std::net::IpAddr::from_str(expected_guest_addr).unwrap() {
94                         s = format!(
95                             "Expecting the guest ip '{}' while being connected with ip '{}'",
96                             expected_guest_addr,
97                             addr.ip()
98                         );
99                         return Err(WaitForBootError::WrongGuestAddr);
100                     }
101 
102                     Ok(())
103                 }
104                 Err(e) => {
105                     s = "TcpListener::accept() failed".to_string();
106                     Err(WaitForBootError::Accept(e))
107                 }
108             }
109         })() {
110             Err(e) => {
111                 let duration = start.elapsed();
112                 eprintln!(
113                     "\n\n==== Start 'wait_vm_boot' (FAILED) ====\n\n\
114                  duration =\"{:?}, timeout = {}s\"\n\
115                  listen_addr=\"{}\"\n\
116                  expected_guest_addr=\"{}\"\n\
117                  message=\"{}\"\n\
118                  error=\"{:?}\"\n\
119                  \n==== End 'wait_vm_boot' outout ====\n\n",
120                     duration, timeout, listen_addr, expected_guest_addr, s, e
121                 );
122 
123                 Err(e)
124             }
125             Ok(_) => Ok(()),
126         }
127     }
128 }
129 
130 pub enum DiskType {
131     OperatingSystem,
132     CloudInit,
133 }
134 
135 pub trait DiskConfig {
136     fn prepare_files(&mut self, tmp_dir: &TempDir, network: &GuestNetworkConfig);
137     fn prepare_cloudinit(&self, tmp_dir: &TempDir, network: &GuestNetworkConfig) -> String;
138     fn disk(&self, disk_type: DiskType) -> Option<String>;
139 }
140 
141 pub struct UbuntuDiskConfig {
142     osdisk_path: String,
143     cloudinit_path: String,
144     image_name: String,
145 }
146 
147 impl UbuntuDiskConfig {
148     pub fn new(image_name: String) -> Self {
149         UbuntuDiskConfig {
150             image_name,
151             osdisk_path: String::new(),
152             cloudinit_path: String::new(),
153         }
154     }
155 }
156 
157 pub struct WindowsDiskConfig {
158     image_name: String,
159     osdisk_path: String,
160     loopback_device: String,
161     windows_snapshot_cow: String,
162     windows_snapshot: String,
163 }
164 
165 impl WindowsDiskConfig {
166     pub fn new(image_name: String) -> Self {
167         WindowsDiskConfig {
168             image_name,
169             osdisk_path: String::new(),
170             loopback_device: String::new(),
171             windows_snapshot_cow: String::new(),
172             windows_snapshot: String::new(),
173         }
174     }
175 }
176 
177 impl Drop for WindowsDiskConfig {
178     fn drop(&mut self) {
179         // dmsetup remove windows-snapshot-1
180         std::process::Command::new("dmsetup")
181             .arg("remove")
182             .arg(self.windows_snapshot.as_str())
183             .output()
184             .expect("Expect removing Windows snapshot with 'dmsetup' to succeed");
185 
186         // dmsetup remove windows-snapshot-cow-1
187         std::process::Command::new("dmsetup")
188             .arg("remove")
189             .arg(self.windows_snapshot_cow.as_str())
190             .output()
191             .expect("Expect removing Windows snapshot CoW with 'dmsetup' to succeed");
192 
193         // losetup -d <loopback_device>
194         std::process::Command::new("losetup")
195             .args(&["-d", self.loopback_device.as_str()])
196             .output()
197             .expect("Expect removing loopback device to succeed");
198     }
199 }
200 
201 impl DiskConfig for UbuntuDiskConfig {
202     fn prepare_cloudinit(&self, tmp_dir: &TempDir, network: &GuestNetworkConfig) -> String {
203         let cloudinit_file_path =
204             String::from(tmp_dir.as_path().join("cloudinit").to_str().unwrap());
205 
206         let cloud_init_directory = tmp_dir.as_path().join("cloud-init").join("ubuntu");
207 
208         fs::create_dir_all(&cloud_init_directory)
209             .expect("Expect creating cloud-init directory to succeed");
210 
211         let source_file_dir = std::env::current_dir()
212             .unwrap()
213             .join("test_data")
214             .join("cloud-init")
215             .join("ubuntu");
216 
217         vec!["meta-data"].iter().for_each(|x| {
218             rate_limited_copy(source_file_dir.join(x), cloud_init_directory.join(x))
219                 .expect("Expect copying cloud-init meta-data to succeed");
220         });
221 
222         let mut user_data_string = String::new();
223         fs::File::open(source_file_dir.join("user-data"))
224             .unwrap()
225             .read_to_string(&mut user_data_string)
226             .expect("Expected reading user-data file in to succeed");
227         user_data_string = user_data_string.replace(
228             "@DEFAULT_TCP_LISTENER_MESSAGE",
229             DEFAULT_TCP_LISTENER_MESSAGE,
230         );
231         user_data_string = user_data_string.replace("@HOST_IP", &network.host_ip);
232         user_data_string =
233             user_data_string.replace("@TCP_LISTENER_PORT", &network.tcp_listener_port.to_string());
234 
235         fs::File::create(cloud_init_directory.join("user-data"))
236             .unwrap()
237             .write_all(user_data_string.as_bytes())
238             .expect("Expected writing out user-data to succeed");
239 
240         let mut network_config_string = String::new();
241 
242         fs::File::open(source_file_dir.join("network-config"))
243             .unwrap()
244             .read_to_string(&mut network_config_string)
245             .expect("Expected reading network-config file in to succeed");
246 
247         network_config_string = network_config_string.replace("192.168.2.1", &network.host_ip);
248         network_config_string = network_config_string.replace("192.168.2.2", &network.guest_ip);
249         network_config_string = network_config_string.replace("192.168.2.3", &network.l2_guest_ip1);
250         network_config_string = network_config_string.replace("192.168.2.4", &network.l2_guest_ip2);
251         network_config_string = network_config_string.replace("192.168.2.5", &network.l2_guest_ip3);
252         network_config_string =
253             network_config_string.replace("12:34:56:78:90:ab", &network.guest_mac);
254         network_config_string =
255             network_config_string.replace("de:ad:be:ef:12:34", &network.l2_guest_mac1);
256         network_config_string =
257             network_config_string.replace("de:ad:be:ef:34:56", &network.l2_guest_mac2);
258         network_config_string =
259             network_config_string.replace("de:ad:be:ef:56:78", &network.l2_guest_mac3);
260 
261         fs::File::create(cloud_init_directory.join("network-config"))
262             .unwrap()
263             .write_all(network_config_string.as_bytes())
264             .expect("Expected writing out network-config to succeed");
265 
266         std::process::Command::new("mkdosfs")
267             .args(&["-n", "cidata"])
268             .args(&["-C", cloudinit_file_path.as_str()])
269             .arg("8192")
270             .output()
271             .expect("Expect creating disk image to succeed");
272 
273         vec!["user-data", "meta-data", "network-config"]
274             .iter()
275             .for_each(|x| {
276                 std::process::Command::new("mcopy")
277                     .arg("-o")
278                     .args(&["-i", cloudinit_file_path.as_str()])
279                     .args(&["-s", cloud_init_directory.join(x).to_str().unwrap(), "::"])
280                     .output()
281                     .expect("Expect copying files to disk image to succeed");
282             });
283 
284         cloudinit_file_path
285     }
286 
287     fn prepare_files(&mut self, tmp_dir: &TempDir, network: &GuestNetworkConfig) {
288         let mut workload_path = dirs::home_dir().unwrap();
289         workload_path.push("workloads");
290 
291         let mut osdisk_base_path = workload_path;
292         osdisk_base_path.push(&self.image_name);
293 
294         let osdisk_path = String::from(tmp_dir.as_path().join("osdisk.img").to_str().unwrap());
295         let cloudinit_path = self.prepare_cloudinit(tmp_dir, network);
296 
297         rate_limited_copy(osdisk_base_path, &osdisk_path)
298             .expect("copying of OS source disk image failed");
299 
300         self.cloudinit_path = cloudinit_path;
301         self.osdisk_path = osdisk_path;
302     }
303 
304     fn disk(&self, disk_type: DiskType) -> Option<String> {
305         match disk_type {
306             DiskType::OperatingSystem => Some(self.osdisk_path.clone()),
307             DiskType::CloudInit => Some(self.cloudinit_path.clone()),
308         }
309     }
310 }
311 
312 impl DiskConfig for WindowsDiskConfig {
313     fn prepare_cloudinit(&self, _tmp_dir: &TempDir, _network: &GuestNetworkConfig) -> String {
314         String::new()
315     }
316 
317     fn prepare_files(&mut self, tmp_dir: &TempDir, _network: &GuestNetworkConfig) {
318         let mut workload_path = dirs::home_dir().unwrap();
319         workload_path.push("workloads");
320 
321         let mut osdisk_path = workload_path;
322         osdisk_path.push(&self.image_name);
323 
324         let osdisk_blk_size = fs::metadata(osdisk_path)
325             .expect("Expect retrieving Windows image metadata")
326             .len()
327             >> 9;
328 
329         let snapshot_cow_path =
330             String::from(tmp_dir.as_path().join("snapshot_cow").to_str().unwrap());
331 
332         // Create and truncate CoW file for device mapper
333         let cow_file_size: u64 = 1 << 30;
334         let cow_file_blk_size = cow_file_size >> 9;
335         let cow_file = std::fs::File::create(snapshot_cow_path.as_str())
336             .expect("Expect creating CoW image to succeed");
337         cow_file
338             .set_len(cow_file_size)
339             .expect("Expect truncating CoW image to succeed");
340 
341         // losetup --find --show /tmp/snapshot_cow
342         let loopback_device = std::process::Command::new("losetup")
343             .arg("--find")
344             .arg("--show")
345             .arg(snapshot_cow_path.as_str())
346             .output()
347             .expect("Expect creating loopback device from snapshot CoW image to succeed");
348 
349         self.loopback_device = String::from_utf8_lossy(&loopback_device.stdout)
350             .trim()
351             .to_string();
352 
353         let random_extension = tmp_dir.as_path().file_name().unwrap();
354         let windows_snapshot_cow = format!(
355             "windows-snapshot-cow-{}",
356             random_extension.to_str().unwrap()
357         );
358 
359         // dmsetup create windows-snapshot-cow-1 --table '0 2097152 linear /dev/loop1 0'
360         std::process::Command::new("dmsetup")
361             .arg("create")
362             .arg(windows_snapshot_cow.as_str())
363             .args(&[
364                 "--table",
365                 format!("0 {} linear {} 0", cow_file_blk_size, self.loopback_device).as_str(),
366             ])
367             .output()
368             .expect("Expect creating Windows snapshot CoW with 'dmsetup' to succeed");
369 
370         let windows_snapshot = format!("windows-snapshot-{}", random_extension.to_str().unwrap());
371 
372         // dmsetup mknodes
373         std::process::Command::new("dmsetup")
374             .arg("mknodes")
375             .output()
376             .expect("Expect device mapper nodes to be ready");
377 
378         // dmsetup create windows-snapshot-1 --table '0 41943040 snapshot /dev/mapper/windows-base /dev/mapper/windows-snapshot-cow-1 P 8'
379         std::process::Command::new("dmsetup")
380             .arg("create")
381             .arg(windows_snapshot.as_str())
382             .args(&[
383                 "--table",
384                 format!(
385                     "0 {} snapshot /dev/mapper/windows-base /dev/mapper/{} P 8",
386                     osdisk_blk_size,
387                     windows_snapshot_cow.as_str()
388                 )
389                 .as_str(),
390             ])
391             .output()
392             .expect("Expect creating Windows snapshot with 'dmsetup' to succeed");
393 
394         // dmsetup mknodes
395         std::process::Command::new("dmsetup")
396             .arg("mknodes")
397             .output()
398             .expect("Expect device mapper nodes to be ready");
399 
400         self.osdisk_path = format!("/dev/mapper/{}", windows_snapshot);
401         self.windows_snapshot_cow = windows_snapshot_cow;
402         self.windows_snapshot = windows_snapshot;
403     }
404 
405     fn disk(&self, disk_type: DiskType) -> Option<String> {
406         match disk_type {
407             DiskType::OperatingSystem => Some(self.osdisk_path.clone()),
408             DiskType::CloudInit => None,
409         }
410     }
411 }
412 
413 pub fn rate_limited_copy<P: AsRef<Path>, Q: AsRef<Path>>(from: P, to: Q) -> io::Result<u64> {
414     for i in 0..10 {
415         let free_bytes = unsafe {
416             let mut stats = std::mem::MaybeUninit::zeroed();
417             let fs_name = std::ffi::CString::new("/tmp").unwrap();
418             libc::statvfs(fs_name.as_ptr(), stats.as_mut_ptr());
419 
420             let free_blocks = stats.assume_init().f_bfree;
421             let block_size = stats.assume_init().f_bsize;
422 
423             free_blocks * block_size
424         };
425 
426         // Make sure there is at least 6 GiB of space
427         if free_bytes < 6 << 30 {
428             eprintln!(
429                 "Not enough space on disk ({}). Attempt {} of 10. Sleeping.",
430                 free_bytes, i
431             );
432             thread::sleep(std::time::Duration::new(60, 0));
433             continue;
434         }
435 
436         match fs::copy(&from, &to) {
437             Err(e) => {
438                 if let Some(errno) = e.raw_os_error() {
439                     if errno == libc::ENOSPC {
440                         eprintln!("Copy returned ENOSPC. Attempt {} of 10. Sleeping.", i);
441                         thread::sleep(std::time::Duration::new(60, 0));
442                         continue;
443                     }
444                 }
445                 return Err(e);
446             }
447             Ok(i) => return Ok(i),
448         }
449     }
450     Err(io::Error::last_os_error())
451 }
452 
453 pub fn handle_child_output(
454     r: Result<(), std::boxed::Box<dyn std::any::Any + std::marker::Send>>,
455     output: &std::process::Output,
456 ) {
457     use std::os::unix::process::ExitStatusExt;
458     if r.is_ok() && output.status.success() {
459         return;
460     }
461 
462     match output.status.code() {
463         None => {
464             // Don't treat child.kill() as a problem
465             if output.status.signal() == Some(9) && r.is_ok() {
466                 return;
467             }
468 
469             eprintln!(
470                 "==== child killed by signal: {} ====",
471                 output.status.signal().unwrap()
472             );
473         }
474         Some(code) => {
475             eprintln!("\n\n==== child exit code: {} ====", code);
476         }
477     }
478 
479     eprintln!(
480         "\n\n==== Start child stdout ====\n\n{}\n\n==== End child stdout ====",
481         String::from_utf8_lossy(&output.stdout)
482     );
483     eprintln!(
484         "\n\n==== Start child stderr ====\n\n{}\n\n==== End child stderr ====",
485         String::from_utf8_lossy(&output.stderr)
486     );
487 
488     panic!("Test failed")
489 }
490 
491 #[derive(Debug)]
492 pub struct PasswordAuth {
493     pub username: String,
494     pub password: String,
495 }
496 
497 pub const DEFAULT_SSH_RETRIES: u8 = 6;
498 pub const DEFAULT_SSH_TIMEOUT: u8 = 10;
499 
500 #[derive(Debug)]
501 pub enum SshCommandError {
502     Connection(std::io::Error),
503     Handshake(ssh2::Error),
504     Authentication(ssh2::Error),
505     ChannelSession(ssh2::Error),
506     Command(ssh2::Error),
507     ExitStatus(ssh2::Error),
508     NonZeroExitStatus(i32),
509 }
510 
511 pub fn ssh_command_ip_with_auth(
512     command: &str,
513     auth: &PasswordAuth,
514     ip: &str,
515     retries: u8,
516     timeout: u8,
517 ) -> Result<String, SshCommandError> {
518     let mut s = String::new();
519 
520     let mut counter = 0;
521     loop {
522         match (|| -> Result<(), SshCommandError> {
523             let tcp =
524                 TcpStream::connect(format!("{}:22", ip)).map_err(SshCommandError::Connection)?;
525             let mut sess = Session::new().unwrap();
526             sess.set_tcp_stream(tcp);
527             sess.handshake().map_err(SshCommandError::Handshake)?;
528 
529             sess.userauth_password(&auth.username, &auth.password)
530                 .map_err(SshCommandError::Authentication)?;
531             assert!(sess.authenticated());
532 
533             let mut channel = sess
534                 .channel_session()
535                 .map_err(SshCommandError::ChannelSession)?;
536             channel.exec(command).map_err(SshCommandError::Command)?;
537 
538             // Intentionally ignore these results here as their failure
539             // does not precipitate a repeat
540             let _ = channel.read_to_string(&mut s);
541             let _ = channel.close();
542             let _ = channel.wait_close();
543 
544             let status = channel.exit_status().map_err(SshCommandError::ExitStatus)?;
545 
546             if status != 0 {
547                 Err(SshCommandError::NonZeroExitStatus(status))
548             } else {
549                 Ok(())
550             }
551         })() {
552             Ok(_) => break,
553             Err(e) => {
554                 counter += 1;
555                 if counter >= retries {
556                     eprintln!(
557                         "\n\n==== Start ssh command output (FAILED) ====\n\n\
558                          command=\"{}\"\n\
559                          auth=\"{:#?}\"\n\
560                          ip=\"{}\"\n\
561                          output=\"{}\"\n\
562                          error=\"{:?}\"\n\
563                          \n==== End ssh command outout ====\n\n",
564                         command, auth, ip, s, e
565                     );
566 
567                     return Err(e);
568                 }
569             }
570         };
571         thread::sleep(std::time::Duration::new((timeout * counter).into(), 0));
572     }
573     Ok(s)
574 }
575 
576 pub fn ssh_command_ip(
577     command: &str,
578     ip: &str,
579     retries: u8,
580     timeout: u8,
581 ) -> Result<String, SshCommandError> {
582     ssh_command_ip_with_auth(
583         command,
584         &PasswordAuth {
585             username: String::from("cloud"),
586             password: String::from("cloud123"),
587         },
588         ip,
589         retries,
590         timeout,
591     )
592 }
593 
594 pub fn exec_host_command_status(command: &str) -> ExitStatus {
595     std::process::Command::new("bash")
596         .args(&["-c", command])
597         .status()
598         .unwrap_or_else(|_| panic!("Expected '{}' to run", command))
599 }
600 
601 pub fn exec_host_command_output(command: &str) -> Output {
602     std::process::Command::new("bash")
603         .args(&["-c", command])
604         .output()
605         .unwrap_or_else(|_| panic!("Expected '{}' to run", command))
606 }
607