1use std::sync::{mpsc, Arc, Mutex};
16
17use base64::engine::general_purpose::STANDARD as B64;
18use base64::Engine as _;
19use serde_json::{json, Value};
20
21use crate::args::{strvec, NetOpts};
22use crate::error::{Error, Result};
23use crate::transport::{http_request, normalize_url, ws_url, WsTransport, TOKEN_ENV, URL_ENV};
24use crate::types::{CommandResult, RemoteExecResult, SandboxInfo, ShellSessionInfo};
25
26const MACHINE_FIELDS: &str = "id name image kind command status running exitCode pid detached \
31 cpus mem volume stateDir createdAt finishedAt network netIp \
32 ports { bind host guest }";
33const CMD_RESULT_FIELDS: &str = "exitCode stdout stderr";
34const SESSION_FIELDS: &str = "id machineId finished truncated";
35
36fn list_query() -> String {
37 format!("query($all: Boolean!) {{ machines(all: $all) {{ {MACHINE_FIELDS} }} }}")
38}
39fn get_query() -> String {
40 format!("query($id: String!) {{ machine(id: $id) {{ {MACHINE_FIELDS} }} }}")
41}
42const LOGS_QUERY: &str =
43 "query($id: String!, $boot: Boolean!) { machineLogs(id: $id, boot: $boot) }";
44
45fn stop_mutation() -> String {
46 format!("mutation($id: String!) {{ stopMachine(id: $id) {{ {CMD_RESULT_FIELDS} }} }}")
47}
48fn start_mutation() -> String {
49 format!("mutation($id: String!) {{ startMachine(id: $id) {{ {CMD_RESULT_FIELDS} }} }}")
50}
51fn remove_mutation() -> String {
52 format!(
53 "mutation($ids: [String!]!, $force: Boolean!) {{ \
54 removeMachines(ids: $ids, force: $force) {{ {CMD_RESULT_FIELDS} }} }}"
55 )
56}
57fn update_mutation() -> String {
58 format!(
59 "mutation($id: String!, $cpus: Int, $mem: Int) {{ \
60 updateMachine(id: $id, cpus: $cpus, mem: $mem) {{ {CMD_RESULT_FIELDS} }} }}"
61 )
62}
63fn commit_mutation() -> String {
64 format!(
65 "mutation($id: String!, $name: String!, $description: String!) {{ \
66 commitMachine(id: $id, name: $name, description: $description) {{ {CMD_RESULT_FIELDS} }} }}"
67 )
68}
69
70const RUN_LINUX_MUTATION: &str = "mutation($input: RunLinuxInput!) { runLinux(input: $input) }";
71const RUN_BSD_MUTATION: &str = "mutation($input: RunBsdInput!) { runBsd(input: $input) }";
72const RUN_NANOS_MUTATION: &str = "mutation($input: RunNanosInput!) { runNanos(input: $input) }";
73const RUN_UNIKRAFT_MUTATION: &str =
74 "mutation($input: RunUnikraftInput!) { runUnikraft(input: $input) }";
75const RUN_SOLO5_MUTATION: &str = "mutation($input: RunSolo5Input!) { runSolo5(input: $input) }";
76const RUN_OSV_MUTATION: &str = "mutation($input: RunOsvInput!) { runOsv(input: $input) }";
77const RUN_FLAVOR_MUTATION: &str = "mutation($input: RunFlavorInput!) { runFlavor(input: $input) }";
78
79const MACHINE_LOGS_SUBSCRIPTION: &str =
80 "subscription($id: String!, $follow: Boolean!, $boot: Boolean!) { \
81 machineLogs(id: $id, follow: $follow, boot: $boot) { dataBase64 exitCode } }";
82
83fn open_shell_mutation() -> String {
84 format!(
85 "mutation($machineId: String!, $command: [String!]!, $env: [String!]!, \
86 $rows: Int!, $cols: Int!) {{ \
87 openShell(machineId: $machineId, command: $command, env: $env, \
88 rows: $rows, cols: $cols) {{ {SESSION_FIELDS} }} }}"
89 )
90}
91const SHELL_OUTPUT_SUBSCRIPTION: &str = "subscription($sessionId: String!) { \
92 shellOutput(sessionId: $sessionId) { dataBase64 exitCode } }";
93const SEND_INPUT_MUTATION: &str = "mutation($sessionId: String!, $dataBase64: String!) { \
94 sendShellInput(sessionId: $sessionId, dataBase64: $dataBase64) }";
95const RESIZE_MUTATION: &str = "mutation($sessionId: String!, $rows: Int!, $cols: Int!) { \
96 resizeShell(sessionId: $sessionId, rows: $rows, cols: $cols) }";
97const CLOSE_MUTATION: &str = "mutation($sessionId: String!) { closeShell(sessionId: $sessionId) }";
98
99struct ClientInner {
104 url: String,
105 token: String,
106 ws: Mutex<Option<Arc<WsTransport>>>,
110}
111
112#[derive(Clone)]
119pub struct Client {
120 inner: Arc<ClientInner>,
121}
122
123impl std::fmt::Debug for Client {
124 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
125 f.debug_struct("Client")
126 .field("url", &self.inner.url)
127 .finish()
128 }
129}
130
131impl Client {
132 pub fn new(url: impl Into<String>, token: impl Into<String>) -> Result<Client> {
137 let url = normalize_url(&url.into());
138 if url.is_empty() {
139 return Err(Error::InvalidInput("the daemon URL is empty".into()));
140 }
141 let token = token.into().trim().to_string();
142 if token.is_empty() {
143 return Err(Error::InvalidInput(
144 "a daemon URL without a token is a configuration error; pass the bearer token"
145 .into(),
146 ));
147 }
148 Ok(Client {
149 inner: Arc::new(ClientInner {
150 url,
151 token,
152 ws: Mutex::new(None),
153 }),
154 })
155 }
156
157 pub fn from_env() -> Result<Client> {
164 let url = std::env::var(URL_ENV)
165 .unwrap_or_default()
166 .trim()
167 .to_string();
168 if url.is_empty() {
169 return Err(Error::InvalidInput(format!(
170 "{URL_ENV} is not set; nothing to connect to"
171 )));
172 }
173 let token = std::env::var(TOKEN_ENV)
174 .unwrap_or_default()
175 .trim()
176 .to_string();
177 if token.is_empty() {
178 return Err(Error::InvalidInput(format!(
179 "{URL_ENV} is set but {TOKEN_ENV} is not"
180 )));
181 }
182 Client::new(url, token)
183 }
184
185 pub fn url(&self) -> &str {
187 &self.inner.url
188 }
189
190 pub fn request(&self, query: &str, variables: Value) -> Result<Value> {
194 http_request(&self.inner.url, &self.inner.token, query, &variables)
195 }
196
197 pub fn subscribe(
200 &self,
201 query: &str,
202 variables: Value,
203 on_next: impl FnMut(Value) + Send + 'static,
204 ) -> Result<Subscription> {
205 self.subscribe_with(query, variables, on_next, |_| {}, || {})
206 }
207
208 pub fn subscribe_with(
210 &self,
211 query: &str,
212 variables: Value,
213 on_next: impl FnMut(Value) + Send + 'static,
214 on_error: impl FnMut(Error) + Send + 'static,
215 on_complete: impl FnMut() + Send + 'static,
216 ) -> Result<Subscription> {
217 let transport = self.ws();
218 let id = transport.subscribe(
219 query,
220 variables,
221 Box::new(on_next),
222 Box::new(on_error),
223 Box::new(on_complete),
224 )?;
225 Ok(Subscription { transport, id })
226 }
227
228 fn ws(&self) -> Arc<WsTransport> {
229 let mut guard = self.inner.ws.lock().unwrap();
230 guard
231 .get_or_insert_with(|| {
232 Arc::new(WsTransport::new(
233 ws_url(&self.inner.url),
234 self.inner.token.clone(),
235 ))
236 })
237 .clone()
238 }
239
240 pub fn list(&self, all: bool) -> Result<Vec<SandboxInfo>> {
244 let data = self.request(&list_query(), json!({"all": all}))?;
245 Ok(data
246 .get("machines")
247 .and_then(Value::as_array)
248 .map(|machines| machines.iter().map(SandboxInfo::from_graphql).collect())
249 .unwrap_or_default())
250 }
251
252 pub fn get(&self, id: &str) -> Result<Option<SandboxInfo>> {
254 let data = self.request(&get_query(), json!({"id": id}))?;
255 Ok(data
256 .get("machine")
257 .filter(|m| !m.is_null())
258 .map(SandboxInfo::from_graphql))
259 }
260
261 pub fn stop(&self, id: &str) -> Result<CommandResult> {
262 let data = self.request(&stop_mutation(), json!({"id": id}))?;
263 Ok(CommandResult::from_graphql(&data["stopMachine"]))
264 }
265
266 pub fn start(&self, id: &str) -> Result<CommandResult> {
267 let data = self.request(&start_mutation(), json!({"id": id}))?;
268 Ok(CommandResult::from_graphql(&data["startMachine"]))
269 }
270
271 pub fn remove<S: AsRef<str>>(&self, ids: &[S], force: bool) -> Result<CommandResult> {
272 let ids: Vec<&str> = ids.iter().map(AsRef::as_ref).collect();
273 let data = self.request(&remove_mutation(), json!({"ids": ids, "force": force}))?;
274 Ok(CommandResult::from_graphql(&data["removeMachines"]))
275 }
276
277 pub fn update(&self, id: &str, cpus: Option<u32>, mem: Option<u32>) -> Result<CommandResult> {
279 let data = self.request(
280 &update_mutation(),
281 json!({"id": id, "cpus": cpus, "mem": mem}),
282 )?;
283 Ok(CommandResult::from_graphql(&data["updateMachine"]))
284 }
285
286 pub fn commit(&self, id: &str, name: &str, description: &str) -> Result<CommandResult> {
288 let data = self.request(
289 &commit_mutation(),
290 json!({"id": id, "name": name, "description": description}),
291 )?;
292 Ok(CommandResult::from_graphql(&data["commitMachine"]))
293 }
294
295 pub fn logs(&self, id: &str, boot: bool) -> Result<String> {
298 let data = self.request(LOGS_QUERY, json!({"id": id, "boot": boot}))?;
299 Ok(data
300 .get("machineLogs")
301 .and_then(Value::as_str)
302 .unwrap_or_default()
303 .to_string())
304 }
305
306 pub fn follow_logs(&self, id: &str) -> FollowLogsBuilder {
317 FollowLogsBuilder {
318 client: self.clone(),
319 id: id.to_string(),
320 follow: true,
321 boot: false,
322 on_data: None,
323 on_error: None,
324 on_complete: None,
325 }
326 }
327
328 pub fn run_linux(&self) -> RunLinuxBuilder {
332 RunLinuxBuilder {
333 client: self.clone(),
334 image: None,
335 cpus: None,
336 mem: None,
337 net: NetOpts::default(),
338 volume: None,
339 mounts: Vec::new(),
340 env: Vec::new(),
341 entrypoint: None,
342 initramfs: false,
343 kernel: None,
344 kernel_version: None,
345 console: None,
346 repo: None,
347 command: Vec::new(),
348 }
349 }
350
351 pub fn run_bsd(&self, os: BsdOs) -> RunBsdBuilder {
353 RunBsdBuilder {
354 client: self.clone(),
355 os,
356 version: None,
357 cpus: None,
358 mem: None,
359 net: NetOpts::default(),
360 volume: None,
361 persist: false,
362 force: false,
363 firmware: None,
364 attach_disk: Vec::new(),
365 disk_size: None,
366 repo: None,
367 command: Vec::new(),
368 }
369 }
370
371 pub fn run_nanos(&self) -> RunNanosBuilder {
373 RunNanosBuilder {
374 client: self.clone(),
375 image: None,
376 cpus: None,
377 mem: None,
378 net: NetOpts::default(),
379 kernel: None,
380 cmdline: None,
381 persist: false,
382 }
383 }
384
385 pub fn run_unikraft(&self) -> RunUnikraftBuilder {
387 RunUnikraftBuilder {
388 client: self.clone(),
389 path: None,
390 cpus: None,
391 mem: None,
392 net: NetOpts::default(),
393 cmdline: None,
394 initramfs: None,
395 mounts: Vec::new(),
396 }
397 }
398
399 pub fn run_solo5(&self) -> RunSolo5Builder {
404 RunSolo5Builder {
405 client: self.clone(),
406 path: None,
407 cpus: None,
408 mem: None,
409 net: NetOpts::default(),
410 block: Vec::new(),
411 args: Vec::new(),
412 }
413 }
414
415 pub fn run_osv(&self) -> RunOsvBuilder {
417 RunOsvBuilder {
418 client: self.clone(),
419 image: None,
420 cpus: None,
421 mem: None,
422 net: NetOpts::default(),
423 cmdline: None,
424 disk: None,
425 no_disk: false,
426 attach_disk: Vec::new(),
427 gic: None,
428 persist: false,
429 volume: None,
430 }
431 }
432
433 pub fn run_flavor(&self, name: impl Into<String>) -> RunFlavorBuilder {
435 RunFlavorBuilder {
436 client: self.clone(),
437 name: name.into(),
438 cpus: None,
439 mem: None,
440 ports: Vec::new(),
441 volume: None,
442 repo: None,
443 }
444 }
445
446 pub fn exec<I, S>(&self, id: &str, command: I) -> Result<RemoteExecResult>
450 where
451 I: IntoIterator<Item = S>,
452 S: Into<String>,
453 {
454 self.exec_with_env(id, command, Vec::<String>::new())
455 }
456
457 pub fn exec_with_env<I, S, E, T>(
468 &self,
469 id: &str,
470 command: I,
471 env: E,
472 ) -> Result<RemoteExecResult>
473 where
474 I: IntoIterator<Item = S>,
475 S: Into<String>,
476 E: IntoIterator<Item = T>,
477 T: Into<String>,
478 {
479 let transport = self.ws();
480 let data = self.request(
481 &open_shell_mutation(),
482 json!({
483 "machineId": id,
484 "command": strvec(command),
485 "env": strvec(env),
486 "rows": 24,
487 "cols": 80,
488 }),
489 )?;
490 let session = ShellSessionInfo::from_graphql(&data["openShell"]);
491
492 let chunks = Arc::new(Mutex::new(Vec::<u8>::new()));
493 let (done_tx, done_rx) = mpsc::channel::<Result<i32>>();
498
499 let chunk_sink = Arc::clone(&chunks);
500 let exit_tx = done_tx.clone();
501 let error_tx = done_tx.clone();
502 let complete_tx = done_tx;
503
504 let outcome: Result<i32> = (|| {
505 let sub_id = transport.subscribe(
506 SHELL_OUTPUT_SUBSCRIPTION,
507 json!({"sessionId": session.id}),
508 Box::new(move |data: Value| {
509 let payload = &data["shellOutput"];
510 if let Some(b64) = payload["dataBase64"].as_str() {
511 if let Ok(bytes) = B64.decode(b64) {
512 chunk_sink.lock().unwrap().extend_from_slice(&bytes);
513 }
514 }
515 if let Some(code) = payload["exitCode"].as_i64() {
516 let _ = exit_tx.send(Ok(code as i32));
517 }
518 }),
519 Box::new(move |err: Error| {
520 let _ = error_tx.send(Err(err));
521 }),
522 Box::new(move || {
523 let _ = complete_tx.send(Err(Error::GraphQL {
527 message: "shell session ended before an exit code arrived".to_string(),
528 code: None,
529 }));
530 }),
531 )?;
532 let outcome = done_rx.recv().unwrap_or_else(|_| {
533 Err(Error::GraphQL {
534 message: "the shell output subscription was dropped".to_string(),
535 code: None,
536 })
537 });
538 transport.unsubscribe(&sub_id);
539 outcome
540 })();
541
542 let _ = self.request(CLOSE_MUTATION, json!({"sessionId": session.id}));
545
546 let exit_code = outcome?;
547 let output = chunks.lock().unwrap().clone();
548 Ok(RemoteExecResult { exit_code, output })
549 }
550
551 pub fn shell(&self, id: &str) -> ShellBuilder {
562 ShellBuilder {
563 client: self.clone(),
564 machine_id: id.to_string(),
565 command: Vec::new(),
566 env: Vec::new(),
567 rows: 24,
568 cols: 80,
569 }
570 }
571}
572
573pub struct Subscription {
577 transport: Arc<WsTransport>,
578 id: String,
579}
580
581impl Subscription {
582 pub fn id(&self) -> &str {
584 &self.id
585 }
586
587 pub fn unsubscribe(self) {
589 self.transport.unsubscribe(&self.id);
590 }
591}
592
593type DataFn = Box<dyn FnMut(Vec<u8>) + Send>;
598type ErrFn = Box<dyn FnMut(Error) + Send>;
599type DoneFn = Box<dyn FnMut() + Send>;
600
601pub struct FollowLogsBuilder {
603 client: Client,
604 id: String,
605 follow: bool,
606 boot: bool,
607 on_data: Option<DataFn>,
608 on_error: Option<ErrFn>,
609 on_complete: Option<DoneFn>,
610}
611
612impl FollowLogsBuilder {
613 pub fn follow(mut self, follow: bool) -> Self {
615 self.follow = follow;
616 self
617 }
618
619 pub fn boot(mut self, boot: bool) -> Self {
621 self.boot = boot;
622 self
623 }
624
625 pub fn on_data(mut self, cb: impl FnMut(Vec<u8>) + Send + 'static) -> Self {
627 self.on_data = Some(Box::new(cb));
628 self
629 }
630
631 pub fn on_error(mut self, cb: impl FnMut(Error) + Send + 'static) -> Self {
633 self.on_error = Some(Box::new(cb));
634 self
635 }
636
637 pub fn on_complete(mut self, cb: impl FnMut() + Send + 'static) -> Self {
639 self.on_complete = Some(Box::new(cb));
640 self
641 }
642
643 pub fn start(self) -> Result<Subscription> {
645 let mut on_data = self.on_data.unwrap_or_else(|| Box::new(|_| {}));
646 let on_error = self.on_error.unwrap_or_else(|| Box::new(|_| {}));
647 let on_complete = self.on_complete.unwrap_or_else(|| Box::new(|| {}));
648 let transport = self.client.ws();
649 let id = transport.subscribe(
650 MACHINE_LOGS_SUBSCRIPTION,
651 json!({"id": self.id, "follow": self.follow, "boot": self.boot}),
652 Box::new(move |data: Value| {
653 if let Some(b64) = data
654 .pointer("/machineLogs/dataBase64")
655 .and_then(Value::as_str)
656 {
657 if let Ok(bytes) = B64.decode(b64) {
658 on_data(bytes);
659 }
660 }
661 }),
665 on_error,
666 on_complete,
667 )?;
668 Ok(Subscription { transport, id })
669 }
670}
671
672#[derive(Debug, Clone, Copy, PartialEq, Eq)]
678pub enum BsdOs {
679 Freebsd,
680 Netbsd,
681}
682
683impl BsdOs {
684 fn graphql(self) -> &'static str {
685 match self {
686 BsdOs::Freebsd => "FREEBSD",
687 BsdOs::Netbsd => "NETBSD",
688 }
689 }
690}
691
692fn net_input(net: &NetOpts) -> Value {
693 if !net.touched {
694 return Value::Null;
695 }
696 json!({
697 "noNet": net.no_net,
698 "ports": net.ports,
699 "mac": net.mac,
700 "network": net.network,
701 "name": net.name,
702 })
703}
704
705fn launch_mutation(client: &Client, mutation: &str, key: &str, input: Value) -> Result<String> {
706 let data = client.request(mutation, json!({"input": input}))?;
707 data.get(key)
708 .and_then(Value::as_str)
709 .map(str::to_string)
710 .ok_or_else(|| Error::GraphQL {
711 message: format!("the daemon's {key} response carried no machine id"),
712 code: None,
713 })
714}
715
716macro_rules! remote_net_vm_setters {
720 () => {
721 pub fn cpus(mut self, cpus: u32) -> Self {
723 self.cpus = Some(cpus);
724 self
725 }
726
727 pub fn mem(mut self, mib: u32) -> Self {
729 self.mem = Some(mib);
730 self
731 }
732
733 pub fn port(mut self, forward: impl Into<String>) -> Self {
735 self.net.touched = true;
736 self.net.ports.push(forward.into());
737 self
738 }
739
740 pub fn forward(self, host: u16, guest: u16) -> Self {
742 self.port(format!("{host}:{guest}"))
743 }
744
745 pub fn mac(mut self, mac: impl Into<String>) -> Self {
747 self.net.touched = true;
748 self.net.mac = Some(mac.into());
749 self
750 }
751
752 pub fn network(mut self, network: impl Into<String>) -> Self {
754 self.net.touched = true;
755 self.net.network = Some(network.into());
756 self
757 }
758
759 pub fn name(mut self, name: impl Into<String>) -> Self {
761 self.net.touched = true;
762 self.net.name = Some(name.into());
763 self
764 }
765
766 pub fn no_net(mut self) -> Self {
768 self.net.touched = true;
769 self.net.no_net = true;
770 self
771 }
772 };
773}
774
775pub struct RunLinuxBuilder {
777 client: Client,
778 image: Option<String>,
779 cpus: Option<u32>,
780 mem: Option<u32>,
781 net: NetOpts,
782 volume: Option<String>,
783 mounts: Vec<String>,
784 env: Vec<String>,
785 entrypoint: Option<String>,
786 initramfs: bool,
787 kernel: Option<String>,
788 kernel_version: Option<String>,
789 console: Option<String>,
790 repo: Option<String>,
791 command: Vec<String>,
792}
793
794impl RunLinuxBuilder {
795 remote_net_vm_setters!();
796
797 pub fn image(mut self, image: impl Into<String>) -> Self {
799 self.image = Some(image.into());
800 self
801 }
802
803 pub fn volume(mut self, name: impl Into<String>) -> Self {
805 self.volume = Some(name.into());
806 self
807 }
808
809 pub fn mount(mut self, mount: impl Into<String>) -> Self {
811 self.mounts.push(mount.into());
812 self
813 }
814
815 pub fn env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
817 self.env.push(format!("{}={}", key.into(), value.into()));
818 self
819 }
820
821 pub fn entrypoint(mut self, entrypoint: impl Into<String>) -> Self {
823 self.entrypoint = Some(entrypoint.into());
824 self
825 }
826
827 pub fn initramfs(mut self) -> Self {
829 self.initramfs = true;
830 self
831 }
832
833 pub fn kernel(mut self, kernel: impl Into<String>) -> Self {
835 self.kernel = Some(kernel.into());
836 self
837 }
838
839 pub fn kernel_version(mut self, version: impl Into<String>) -> Self {
841 self.kernel_version = Some(version.into());
842 self
843 }
844
845 pub fn console(mut self, console: impl Into<String>) -> Self {
847 self.console = Some(console.into());
848 self
849 }
850
851 pub fn repo(mut self, repo: impl Into<String>) -> Self {
853 self.repo = Some(repo.into());
854 self
855 }
856
857 pub fn command<I, S>(mut self, command: I) -> Self
859 where
860 I: IntoIterator<Item = S>,
861 S: Into<String>,
862 {
863 self.command = strvec(command);
864 self
865 }
866
867 pub fn launch(self) -> Result<String> {
869 let Some(image) = self.image else {
870 return Err(Error::InvalidInput("run_linux requires an image".into()));
871 };
872 let input = json!({
873 "image": image,
874 "cpus": self.cpus,
875 "mem": self.mem,
876 "net": net_input(&self.net),
877 "volume": self.volume,
878 "mounts": self.mounts,
879 "env": self.env,
880 "entrypoint": self.entrypoint,
881 "initramfs": self.initramfs,
882 "kernel": self.kernel,
883 "kernelVersion": self.kernel_version,
884 "console": self.console,
885 "repo": self.repo,
886 "command": self.command,
887 });
888 launch_mutation(&self.client, RUN_LINUX_MUTATION, "runLinux", input)
889 }
890}
891
892pub struct RunBsdBuilder {
894 client: Client,
895 os: BsdOs,
896 version: Option<String>,
897 cpus: Option<u32>,
898 mem: Option<u32>,
899 net: NetOpts,
900 volume: Option<String>,
901 persist: bool,
902 force: bool,
903 firmware: Option<String>,
904 attach_disk: Vec<String>,
905 disk_size: Option<String>,
906 repo: Option<String>,
907 command: Vec<String>,
908}
909
910impl RunBsdBuilder {
911 remote_net_vm_setters!();
912
913 pub fn version(mut self, version: impl Into<String>) -> Self {
915 self.version = Some(version.into());
916 self
917 }
918
919 pub fn volume(mut self, name: impl Into<String>) -> Self {
921 self.volume = Some(name.into());
922 self
923 }
924
925 pub fn persist(mut self) -> Self {
927 self.persist = true;
928 self
929 }
930
931 pub fn force(mut self) -> Self {
933 self.force = true;
934 self
935 }
936
937 pub fn firmware(mut self, firmware: impl Into<String>) -> Self {
939 self.firmware = Some(firmware.into());
940 self
941 }
942
943 pub fn attach_disk(mut self, disk: impl Into<String>) -> Self {
945 self.attach_disk.push(disk.into());
946 self
947 }
948
949 pub fn disk_size(mut self, size: impl Into<String>) -> Self {
951 self.disk_size = Some(size.into());
952 self
953 }
954
955 pub fn repo(mut self, repo: impl Into<String>) -> Self {
957 self.repo = Some(repo.into());
958 self
959 }
960
961 pub fn command<I, S>(mut self, command: I) -> Self
963 where
964 I: IntoIterator<Item = S>,
965 S: Into<String>,
966 {
967 self.command = strvec(command);
968 self
969 }
970
971 pub fn launch(self) -> Result<String> {
973 let input = json!({
974 "os": self.os.graphql(),
975 "version": self.version,
976 "cpus": self.cpus,
977 "mem": self.mem,
978 "net": net_input(&self.net),
979 "volume": self.volume,
980 "persist": self.persist,
981 "force": self.force,
982 "firmware": self.firmware,
983 "attachDisk": self.attach_disk,
984 "diskSize": self.disk_size,
985 "repo": self.repo,
986 "command": self.command,
987 });
988 launch_mutation(&self.client, RUN_BSD_MUTATION, "runBsd", input)
989 }
990}
991
992pub struct RunNanosBuilder {
994 client: Client,
995 image: Option<String>,
996 cpus: Option<u32>,
997 mem: Option<u32>,
998 net: NetOpts,
999 kernel: Option<String>,
1000 cmdline: Option<String>,
1001 persist: bool,
1002}
1003
1004impl RunNanosBuilder {
1005 remote_net_vm_setters!();
1006
1007 pub fn image(mut self, image: impl Into<String>) -> Self {
1009 self.image = Some(image.into());
1010 self
1011 }
1012
1013 pub fn kernel(mut self, kernel: impl Into<String>) -> Self {
1015 self.kernel = Some(kernel.into());
1016 self
1017 }
1018
1019 pub fn cmdline(mut self, cmdline: impl Into<String>) -> Self {
1021 self.cmdline = Some(cmdline.into());
1022 self
1023 }
1024
1025 pub fn persist(mut self) -> Self {
1027 self.persist = true;
1028 self
1029 }
1030
1031 pub fn launch(self) -> Result<String> {
1033 let Some(image) = self.image else {
1034 return Err(Error::InvalidInput("run_nanos requires an image".into()));
1035 };
1036 let input = json!({
1037 "image": image,
1038 "cpus": self.cpus,
1039 "mem": self.mem,
1040 "net": net_input(&self.net),
1041 "kernel": self.kernel,
1042 "cmdline": self.cmdline,
1043 "persist": self.persist,
1044 });
1045 launch_mutation(&self.client, RUN_NANOS_MUTATION, "runNanos", input)
1046 }
1047}
1048
1049pub struct RunUnikraftBuilder {
1051 client: Client,
1052 path: Option<String>,
1053 cpus: Option<u32>,
1054 mem: Option<u32>,
1055 net: NetOpts,
1056 cmdline: Option<String>,
1057 initramfs: Option<String>,
1058 mounts: Vec<String>,
1059}
1060
1061impl RunUnikraftBuilder {
1062 remote_net_vm_setters!();
1063
1064 pub fn path(mut self, path: impl Into<String>) -> Self {
1066 self.path = Some(path.into());
1067 self
1068 }
1069
1070 pub fn cmdline(mut self, cmdline: impl Into<String>) -> Self {
1072 self.cmdline = Some(cmdline.into());
1073 self
1074 }
1075
1076 pub fn initramfs(mut self, path: impl Into<String>) -> Self {
1078 self.initramfs = Some(path.into());
1079 self
1080 }
1081
1082 pub fn mount(mut self, mount: impl Into<String>) -> Self {
1085 self.mounts.push(mount.into());
1086 self
1087 }
1088
1089 pub fn launch(self) -> Result<String> {
1091 let input = json!({
1092 "path": self.path,
1093 "cpus": self.cpus,
1094 "mem": self.mem,
1095 "net": net_input(&self.net),
1096 "cmdline": self.cmdline,
1097 "initramfs": self.initramfs,
1098 "mounts": self.mounts,
1099 });
1100 launch_mutation(&self.client, RUN_UNIKRAFT_MUTATION, "runUnikraft", input)
1101 }
1102}
1103
1104pub struct RunSolo5Builder {
1106 client: Client,
1107 path: Option<String>,
1108 cpus: Option<u32>,
1109 mem: Option<u32>,
1110 net: NetOpts,
1111 block: Vec<String>,
1112 args: Vec<String>,
1113}
1114
1115impl RunSolo5Builder {
1116 remote_net_vm_setters!();
1117
1118 pub fn path(mut self, path: impl Into<String>) -> Self {
1121 self.path = Some(path.into());
1122 self
1123 }
1124
1125 pub fn block(mut self, block: impl Into<String>) -> Self {
1127 self.block.push(block.into());
1128 self
1129 }
1130
1131 pub fn args<I, S>(mut self, args: I) -> Self
1133 where
1134 I: IntoIterator<Item = S>,
1135 S: Into<String>,
1136 {
1137 self.args = strvec(args);
1138 self
1139 }
1140
1141 pub fn launch(self) -> Result<String> {
1143 let input = json!({
1144 "path": self.path,
1145 "cpus": self.cpus,
1146 "mem": self.mem,
1147 "net": net_input(&self.net),
1148 "block": self.block,
1149 "args": self.args,
1150 });
1151 launch_mutation(&self.client, RUN_SOLO5_MUTATION, "runSolo5", input)
1152 }
1153}
1154
1155pub struct RunOsvBuilder {
1157 client: Client,
1158 image: Option<String>,
1159 cpus: Option<u32>,
1160 mem: Option<u32>,
1161 net: NetOpts,
1162 cmdline: Option<String>,
1163 disk: Option<String>,
1164 no_disk: bool,
1165 attach_disk: Vec<String>,
1166 gic: Option<String>,
1167 persist: bool,
1168 volume: Option<String>,
1169}
1170
1171impl RunOsvBuilder {
1172 remote_net_vm_setters!();
1173
1174 pub fn image(mut self, image: impl Into<String>) -> Self {
1176 self.image = Some(image.into());
1177 self
1178 }
1179
1180 pub fn cmdline(mut self, cmdline: impl Into<String>) -> Self {
1182 self.cmdline = Some(cmdline.into());
1183 self
1184 }
1185
1186 pub fn disk(mut self, disk: impl Into<String>) -> Self {
1188 self.disk = Some(disk.into());
1189 self
1190 }
1191
1192 pub fn no_disk(mut self) -> Self {
1194 self.no_disk = true;
1195 self
1196 }
1197
1198 pub fn attach_disk(mut self, disk: impl Into<String>) -> Self {
1200 self.attach_disk.push(disk.into());
1201 self
1202 }
1203
1204 pub fn gic(mut self, gic: impl Into<String>) -> Self {
1206 self.gic = Some(gic.into());
1207 self
1208 }
1209
1210 pub fn persist(mut self) -> Self {
1212 self.persist = true;
1213 self
1214 }
1215
1216 pub fn volume(mut self, name: impl Into<String>) -> Self {
1218 self.volume = Some(name.into());
1219 self
1220 }
1221
1222 pub fn launch(self) -> Result<String> {
1224 let Some(image) = self.image else {
1225 return Err(Error::InvalidInput("run_osv requires an image".into()));
1226 };
1227 let input = json!({
1228 "image": image,
1229 "cpus": self.cpus,
1230 "mem": self.mem,
1231 "net": net_input(&self.net),
1232 "cmdline": self.cmdline,
1233 "disk": self.disk,
1234 "noDisk": self.no_disk,
1235 "attachDisk": self.attach_disk,
1236 "gic": self.gic,
1237 "persist": self.persist,
1238 "volume": self.volume,
1239 });
1240 launch_mutation(&self.client, RUN_OSV_MUTATION, "runOsv", input)
1241 }
1242}
1243
1244pub struct RunFlavorBuilder {
1246 client: Client,
1247 name: String,
1248 cpus: Option<u32>,
1249 mem: Option<u32>,
1250 ports: Vec<String>,
1251 volume: Option<String>,
1252 repo: Option<String>,
1253}
1254
1255impl RunFlavorBuilder {
1256 pub fn cpus(mut self, cpus: u32) -> Self {
1258 self.cpus = Some(cpus);
1259 self
1260 }
1261
1262 pub fn mem(mut self, mib: u32) -> Self {
1264 self.mem = Some(mib);
1265 self
1266 }
1267
1268 pub fn port(mut self, forward: impl Into<String>) -> Self {
1270 self.ports.push(forward.into());
1271 self
1272 }
1273
1274 pub fn volume(mut self, name: impl Into<String>) -> Self {
1276 self.volume = Some(name.into());
1277 self
1278 }
1279
1280 pub fn repo(mut self, repo: impl Into<String>) -> Self {
1282 self.repo = Some(repo.into());
1283 self
1284 }
1285
1286 pub fn launch(self) -> Result<String> {
1288 let input = json!({
1289 "name": self.name,
1290 "cpus": self.cpus,
1291 "mem": self.mem,
1292 "ports": self.ports,
1293 "volume": self.volume,
1294 "repo": self.repo,
1295 });
1296 launch_mutation(&self.client, RUN_FLAVOR_MUTATION, "runFlavor", input)
1297 }
1298}
1299
1300pub struct ShellBuilder {
1306 client: Client,
1307 machine_id: String,
1308 command: Vec<String>,
1309 env: Vec<String>,
1310 rows: u32,
1311 cols: u32,
1312}
1313
1314impl ShellBuilder {
1315 pub fn command<I, S>(mut self, command: I) -> Self
1317 where
1318 I: IntoIterator<Item = S>,
1319 S: Into<String>,
1320 {
1321 self.command = strvec(command);
1322 self
1323 }
1324
1325 pub fn env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
1327 self.env.push(format!("{}={}", key.into(), value.into()));
1328 self
1329 }
1330
1331 pub fn rows(mut self, rows: u32) -> Self {
1333 self.rows = rows;
1334 self
1335 }
1336
1337 pub fn cols(mut self, cols: u32) -> Self {
1339 self.cols = cols;
1340 self
1341 }
1342
1343 pub fn open(self) -> Result<ShellSession> {
1345 let data = self.client.request(
1346 &open_shell_mutation(),
1347 json!({
1348 "machineId": self.machine_id,
1349 "command": self.command,
1350 "env": self.env,
1351 "rows": self.rows,
1352 "cols": self.cols,
1353 }),
1354 )?;
1355 let info = ShellSessionInfo::from_graphql(&data["openShell"]);
1356 ShellSession::start(self.client, info.id)
1357 }
1358}
1359
1360type OutputFn = Box<dyn FnMut(&[u8]) + Send>;
1361type ExitFn = Box<dyn FnMut(i32) + Send>;
1362
1363struct ShellShared {
1364 output_cb: Option<OutputFn>,
1365 exit_cb: Option<ExitFn>,
1366 buffered_output: Vec<Vec<u8>>,
1372 buffered_exit: Option<i32>,
1373 exit_fired: bool,
1374}
1375
1376pub struct ShellSession {
1384 id: String,
1385 client: Client,
1386 transport: Arc<WsTransport>,
1387 sub_id: String,
1388 shared: Arc<Mutex<ShellShared>>,
1389 closed: bool,
1390}
1391
1392impl std::fmt::Debug for ShellSession {
1393 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1394 f.debug_struct("ShellSession")
1395 .field("id", &self.id)
1396 .finish()
1397 }
1398}
1399
1400impl ShellSession {
1401 fn start(client: Client, session_id: String) -> Result<ShellSession> {
1402 let shared = Arc::new(Mutex::new(ShellShared {
1403 output_cb: None,
1404 exit_cb: None,
1405 buffered_output: Vec::new(),
1406 buffered_exit: None,
1407 exit_fired: false,
1408 }));
1409 let transport = client.ws();
1410
1411 let on_next_shared = Arc::clone(&shared);
1412 let on_error_shared = Arc::clone(&shared);
1413 let sub_id = transport.subscribe(
1414 SHELL_OUTPUT_SUBSCRIPTION,
1415 json!({"sessionId": session_id}),
1416 Box::new(move |data: Value| {
1417 let payload = &data["shellOutput"];
1418 if let Some(b64) = payload["dataBase64"].as_str() {
1419 if let Ok(bytes) = B64.decode(b64) {
1420 emit_output(&on_next_shared, bytes);
1421 }
1422 }
1423 if let Some(code) = payload["exitCode"].as_i64() {
1424 emit_exit(&on_next_shared, code as i32);
1425 }
1426 }),
1427 Box::new(move |_err: Error| {
1428 emit_exit(&on_error_shared, -1);
1433 }),
1434 Box::new(|| {}),
1435 )?;
1436
1437 Ok(ShellSession {
1438 id: session_id,
1439 client,
1440 transport,
1441 sub_id,
1442 shared,
1443 closed: false,
1444 })
1445 }
1446
1447 pub fn id(&self) -> &str {
1449 &self.id
1450 }
1451
1452 pub fn on_output(&self, mut cb: impl FnMut(&[u8]) + Send + 'static) {
1455 let mut shared = self.shared.lock().unwrap();
1456 for chunk in std::mem::take(&mut shared.buffered_output) {
1457 cb(&chunk);
1458 }
1459 shared.output_cb = Some(Box::new(cb));
1460 }
1461
1462 pub fn on_exit(&self, mut cb: impl FnMut(i32) + Send + 'static) {
1464 let mut shared = self.shared.lock().unwrap();
1465 if let Some(code) = shared.buffered_exit.take() {
1466 cb(code);
1467 }
1468 shared.exit_cb = Some(Box::new(cb));
1469 }
1470
1471 pub fn write(&self, data: impl AsRef<[u8]>) -> Result<()> {
1473 self.client.request(
1474 SEND_INPUT_MUTATION,
1475 json!({"sessionId": self.id, "dataBase64": B64.encode(data.as_ref())}),
1476 )?;
1477 Ok(())
1478 }
1479
1480 pub fn resize(&self, rows: u32, cols: u32) -> Result<()> {
1482 self.client.request(
1483 RESIZE_MUTATION,
1484 json!({"sessionId": self.id, "rows": rows, "cols": cols}),
1485 )?;
1486 Ok(())
1487 }
1488
1489 pub fn close(&mut self) {
1491 if self.closed {
1492 return;
1493 }
1494 self.closed = true;
1495 self.transport.unsubscribe(&self.sub_id);
1496 let _ = self
1498 .client
1499 .request(CLOSE_MUTATION, json!({"sessionId": self.id}));
1500 }
1501}
1502
1503fn emit_output(shared: &Arc<Mutex<ShellShared>>, data: Vec<u8>) {
1504 let mut guard = shared.lock().unwrap();
1505 match &mut guard.output_cb {
1506 Some(cb) => cb(&data),
1507 None => guard.buffered_output.push(data),
1508 }
1509}
1510
1511fn emit_exit(shared: &Arc<Mutex<ShellShared>>, code: i32) {
1512 let mut guard = shared.lock().unwrap();
1513 if guard.exit_fired {
1514 return;
1515 }
1516 guard.exit_fired = true;
1517 match &mut guard.exit_cb {
1518 Some(cb) => cb(code),
1519 None => guard.buffered_exit = Some(code),
1520 }
1521}