1use std::collections::HashMap;
2
3use anyhow::Result;
4use clap::{Parser, ValueEnum};
5use krata::{
6 events::EventStream,
7 v1::{
8 common::{
9 zone_image_spec::Image, OciImageFormat, ZoneImageSpec, ZoneKernelOptionsSpec,
10 ZoneOciImageSpec, ZoneResourceSpec, ZoneSpec, ZoneSpecDevice, ZoneState, ZoneTaskSpec,
11 ZoneTaskSpecEnvVar,
12 },
13 control::{
14 control_service_client::ControlServiceClient, watch_events_reply::Event,
15 CreateZoneRequest, PullImageRequest,
16 },
17 },
18};
19use log::error;
20use tokio::select;
21use tonic::{transport::Channel, Request};
22
23use crate::{console::StdioConsoleStream, pull::pull_interactive_progress};
24
25#[derive(ValueEnum, Clone, Debug, PartialEq, Eq)]
26pub enum LaunchImageFormat {
27 Squashfs,
28 Erofs,
29}
30
31#[derive(Parser)]
32#[command(about = "Launch a new zone")]
33pub struct ZoneLaunchCommand {
34 #[arg(long, default_value = "squashfs", help = "Image format")]
35 image_format: LaunchImageFormat,
36 #[arg(long, help = "Overwrite image cache on pull")]
37 pull_overwrite_cache: bool,
38 #[arg(long, help = "Update image on pull")]
39 pull_update: bool,
40 #[arg(short, long, help = "Name of the zone")]
41 name: Option<String>,
42 #[arg(
43 short = 'C',
44 long = "max-cpus",
45 default_value_t = 4,
46 help = "Maximum vCPUs available for the zone"
47 )]
48 max_cpus: u32,
49 #[arg(
50 short = 'c',
51 long = "target-cpus",
52 default_value_t = 1,
53 help = "Target vCPUs for the zone to use"
54 )]
55 target_cpus: u32,
56 #[arg(
57 short = 'M',
58 long = "max-memory",
59 default_value_t = 1024,
60 help = "Maximum memory available to the zone, in megabytes"
61 )]
62 max_memory: u64,
63 #[arg(
64 short = 'm',
65 long = "target-memory",
66 default_value_t = 1024,
67 help = "Target memory for the zone to use, in megabytes"
68 )]
69 target_memory: u64,
70 #[arg[short = 'D', long = "device", help = "Devices to request for the zone"]]
71 device: Vec<String>,
72 #[arg[short, long, help = "Environment variables set in the zone"]]
73 env: Option<Vec<String>>,
74 #[arg(short = 't', long, help = "Allocate tty for task")]
75 tty: bool,
76 #[arg(
77 short,
78 long,
79 help = "Attach to the zone after zone starts, implies --wait"
80 )]
81 attach: bool,
82 #[arg(
83 short = 'W',
84 long,
85 help = "Wait for the zone to start, implied by --attach"
86 )]
87 wait: bool,
88 #[arg(short = 'k', long, help = "OCI kernel image for zone to use")]
89 kernel: Option<String>,
90 #[arg(short = 'I', long, help = "OCI initrd image for zone to use")]
91 initrd: Option<String>,
92 #[arg(short = 'w', long, help = "Working directory")]
93 working_directory: Option<String>,
94 #[arg(long, help = "Enable verbose logging on the kernel")]
95 kernel_verbose: bool,
96 #[arg(long, help = "Additional kernel cmdline options")]
97 kernel_cmdline_append: Option<String>,
98 #[arg(help = "Container image for zone to use")]
99 oci: String,
100 #[arg(
101 allow_hyphen_values = true,
102 trailing_var_arg = true,
103 help = "Command to run inside the zone"
104 )]
105 command: Vec<String>,
106}
107
108impl ZoneLaunchCommand {
109 pub async fn run(
110 self,
111 mut client: ControlServiceClient<Channel>,
112 events: EventStream,
113 ) -> Result<()> {
114 let image = self
115 .pull_image(
116 &mut client,
117 &self.oci,
118 match self.image_format {
119 LaunchImageFormat::Squashfs => OciImageFormat::Squashfs,
120 LaunchImageFormat::Erofs => OciImageFormat::Erofs,
121 },
122 )
123 .await?;
124
125 let kernel = if let Some(ref kernel) = self.kernel {
126 let kernel_image = self
127 .pull_image(&mut client, kernel, OciImageFormat::Tar)
128 .await?;
129 Some(kernel_image)
130 } else {
131 None
132 };
133
134 let initrd = if let Some(ref initrd) = self.initrd {
135 let kernel_image = self
136 .pull_image(&mut client, initrd, OciImageFormat::Tar)
137 .await?;
138 Some(kernel_image)
139 } else {
140 None
141 };
142
143 let request = CreateZoneRequest {
144 spec: Some(ZoneSpec {
145 name: self.name.unwrap_or_default(),
146 image: Some(image),
147 kernel,
148 initrd,
149 initial_resources: Some(ZoneResourceSpec {
150 max_memory: self.max_memory,
151 target_memory: self.target_memory,
152 max_cpus: self.max_cpus,
153 target_cpus: self.target_cpus,
154 }),
155 task: Some(ZoneTaskSpec {
156 environment: env_map(&self.env.unwrap_or_default())
157 .iter()
158 .map(|(key, value)| ZoneTaskSpecEnvVar {
159 key: key.clone(),
160 value: value.clone(),
161 })
162 .collect(),
163 command: self.command,
164 working_directory: self.working_directory.unwrap_or_default(),
165 tty: self.tty,
166 }),
167 annotations: vec![],
168 devices: self
169 .device
170 .iter()
171 .map(|name| ZoneSpecDevice { name: name.clone() })
172 .collect(),
173 kernel_options: Some(ZoneKernelOptionsSpec {
174 verbose: self.kernel_verbose,
175 cmdline_append: self.kernel_cmdline_append.clone().unwrap_or_default(),
176 }),
177 }),
178 };
179 let response = client
180 .create_zone(Request::new(request))
181 .await?
182 .into_inner();
183 let id = response.zone_id;
184
185 if self.wait || self.attach {
186 wait_zone_started(&id, events.clone()).await?;
187 }
188
189 let code = if self.attach {
190 let input = StdioConsoleStream::stdin_stream(id.clone(), true).await;
191 let output = client.attach_zone_console(input).await?.into_inner();
192 let stdout_handle =
193 tokio::task::spawn(async move { StdioConsoleStream::stdout(output, true).await });
194 let exit_hook_task = StdioConsoleStream::zone_exit_hook(id.clone(), events).await?;
195 select! {
196 x = stdout_handle => {
197 x??;
198 None
199 },
200 x = exit_hook_task => x?
201 }
202 } else {
203 println!("{}", id);
204 None
205 };
206 StdioConsoleStream::restore_terminal_mode();
207 std::process::exit(code.unwrap_or(0));
208 }
209
210 async fn pull_image(
211 &self,
212 client: &mut ControlServiceClient<Channel>,
213 image: &str,
214 format: OciImageFormat,
215 ) -> Result<ZoneImageSpec> {
216 let response = client
217 .pull_image(PullImageRequest {
218 image: image.to_string(),
219 format: format.into(),
220 overwrite_cache: self.pull_overwrite_cache,
221 update: self.pull_update,
222 })
223 .await?;
224 let reply = pull_interactive_progress(response.into_inner()).await?;
225 Ok(ZoneImageSpec {
226 image: Some(Image::Oci(ZoneOciImageSpec {
227 digest: reply.digest,
228 format: reply.format,
229 })),
230 })
231 }
232}
233
234async fn wait_zone_started(id: &str, events: EventStream) -> Result<()> {
235 let mut stream = events.subscribe();
236 while let Ok(event) = stream.recv().await {
237 match event {
238 Event::ZoneChanged(changed) => {
239 let Some(zone) = changed.zone else {
240 continue;
241 };
242
243 if zone.id != id {
244 continue;
245 }
246
247 let Some(status) = zone.status else {
248 continue;
249 };
250
251 if let Some(ref error) = status.error_status {
252 if status.state() == ZoneState::Failed {
253 error!("launch failed: {}", error.message);
254 std::process::exit(1);
255 } else {
256 error!("zone error: {}", error.message);
257 }
258 }
259
260 if status.state() == ZoneState::Destroyed {
261 error!("zone destroyed");
262 std::process::exit(1);
263 }
264
265 if status.state() == ZoneState::Created {
266 break;
267 }
268 }
269 }
270 }
271 Ok(())
272}
273
274fn env_map(env: &[String]) -> HashMap<String, String> {
275 let mut map = HashMap::<String, String>::new();
276 for item in env {
277 if let Some((key, value)) = item.split_once('=') {
278 map.insert(key.to_string(), value.to_string());
279 }
280 }
281 map
282}