alice.pds.demo.boltless.dev
This repository has no description
1use crate::cache::{CacheUploadManager, ReadCacheProxy, WriteCacheProxy};
2use crate::command::Spec;
3use crate::dns_proxy::DnsProxy;
4use crate::exec;
5use crate::nix_config::{self, SYSTEMCTL_EXECUTABLE};
6use crate::on_payload;
7use crate::protocol::{self, Message, v1};
8use crate::{activation, command};
9use anyhow::{Context, Result, bail};
10use std::time::Duration;
11use tokio::io::{AsyncWrite, BufReader};
12use tokio::sync::mpsc::{self, Sender};
13use tokio::task::{JoinError, JoinSet};
14use tokio_vsock::{VsockAddr, VsockStream};
15use tracing::{info, warn};
16
17pub async fn run(host_cid: u32, port: u32) -> Result<()> {
18 let mut conn = VsockStream::connect(VsockAddr::new(host_cid, port))
19 .await
20 .with_context(|| format!("dial host vsock cid={host_cid} port={port}"))?;
21
22 send_hello(&mut conn).await?;
23
24 let (reader_conn, writer_conn) = tokio::io::split(conn);
25 let (out_tx, out_rx) = mpsc::channel::<Message>(256);
26 let writer = tokio::spawn(async move { writer_loop(writer_conn, out_rx).await });
27 let mut reader = BufReader::new(reader_conn);
28
29 let init = match protocol::read_message(&mut reader).await? {
30 Some(Message {
31 init: Some(init), ..
32 }) => init,
33 Some(other) => bail!("expected init, got {}", protocol::kind(&other)),
34 None => bail!("read init: EOF"),
35 };
36 info!(job_id = %init.job_id, "received init");
37
38 let read_proxy = ReadCacheProxy::start(host_cid, init.cache_read_proxy_port)
39 .await
40 .context("start read cache proxy")?;
41 let write_proxy = WriteCacheProxy::start(host_cid, init.cache_upload_proxy_port)
42 .await
43 .context("start write cache proxy")?;
44 let _dns_proxy = DnsProxy::start(host_cid, init.dns_proxy_port)
45 .await
46 .context("start dns proxy")?;
47 let _cache_cfg = nix_config::configure(
48 &init,
49 read_proxy.as_ref().map(ReadCacheProxy::url).unwrap_or(""),
50 )
51 .await
52 .context("configure nix cache")?;
53 let uploader = CacheUploadManager::start(
54 write_proxy.as_ref().map(WriteCacheProxy::url).unwrap_or(""),
55 out_tx.clone(),
56 )
57 .await
58 .context("start cache upload manager")?;
59
60 let mut tasks = JoinSet::new();
61 let read_result: Result<()> = loop {
62 tokio::select! {
63 read = protocol::read_message(&mut reader) => match read {
64 Ok(Some(msg)) => spawn_message_task(&mut tasks, host_cid, msg, &out_tx, uploader.clone()),
65 Ok(None) => break Ok(()),
66 Err(error) => break Err(error).context("read message"),
67 },
68 Some(result) = tasks.join_next(), if !tasks.is_empty() => {
69 log_task_result(result, false);
70 }
71 }
72 };
73
74 tasks.abort_all();
75 while let Some(result) = tasks.join_next().await {
76 log_task_result(result, true);
77 }
78
79 drop(out_tx);
80 let _ = writer.await;
81 read_result?;
82 Ok(())
83}
84
85fn spawn_message_task(
86 tasks: &mut JoinSet<()>,
87 host_cid: u32,
88 msg: Message,
89 out_tx: &Sender<Message>,
90 uploader: Option<CacheUploadManager>,
91) {
92 let kind = protocol::kind(&msg);
93 let handle = on_payload!(msg, {
94 activate_config => tasks.spawn(activation::run(msg.id, activate_config, out_tx.clone())),
95 exec_start => tasks.spawn(exec::run(msg.id, exec_start, out_tx.clone(), host_cid)),
96 cache_drain => tasks.spawn(run_cache_drain(msg.id, cache_drain, out_tx.clone(), uploader)),
97 poweroff => tasks.spawn(run_poweroff(msg.id, poweroff, out_tx.clone())),
98 });
99 if handle.is_none() {
100 warn!(kind, "ignoring unsupported message");
101 }
102}
103
104fn log_task_result(result: Result<(), JoinError>, shutting_down: bool) {
105 match result {
106 Ok(()) => {}
107 Err(error) if shutting_down && error.is_cancelled() => {}
108 Err(error) => warn!(%error, "session handler task failed"),
109 }
110}
111
112async fn writer_loop<W>(mut conn: W, mut rx: mpsc::Receiver<Message>)
113where
114 W: AsyncWrite + Unpin,
115{
116 while let Some(msg) = rx.recv().await {
117 if let Err(error) = protocol::write_message(&mut conn, &msg).await {
118 warn!(%error, "failed to write protocol message");
119 break;
120 }
121 }
122}
123
124async fn send_hello(conn: &mut VsockStream) -> Result<()> {
125 let boot_id = tokio::fs::read_to_string("/proc/sys/kernel/random/boot_id")
126 .await
127 .unwrap_or_default()
128 .trim()
129 .to_owned();
130 let nix_version = nix_config::nix_version().await;
131
132 let hello_payload = v1::Hello {
133 protocol_version: protocol::PROTOCOL_VERSION,
134 agent_version: env!("CARGO_PKG_VERSION").to_string(),
135 boot_id: boot_id.clone(),
136 nix_version: nix_version.clone(),
137 };
138 info!(
139 protocol = hello_payload.protocol_version,
140 version = %hello_payload.agent_version,
141 boot = %hello_payload.boot_id,
142 nix = %hello_payload.nix_version,
143 "sent hello"
144 );
145 let hello = Message {
146 id: "hello".to_owned(),
147 hello: Some(hello_payload),
148 ..Default::default()
149 };
150
151 protocol::write_message(conn, &hello)
152 .await
153 .context("send hello")?;
154 Ok(())
155}
156
157async fn run_cache_drain(
158 id: String,
159 req: v1::CacheDrain,
160 out: Sender<Message>,
161 uploader: Option<CacheUploadManager>,
162) {
163 let timeout =
164 (req.timeout_seconds > 0).then(|| Duration::from_secs(u64::from(req.timeout_seconds)));
165 let stats = match uploader.as_ref() {
166 Some(uploader) => uploader.drain(timeout).await,
167 None => Default::default(),
168 };
169
170 if let Some(error) = &stats.last_error {
171 warn!(
172 %id,
173 pending = stats.pending,
174 active = stats.active,
175 uploaded = stats.uploaded,
176 failed = stats.failed,
177 %error,
178 "cache drain completed with error"
179 );
180 } else {
181 info!(
182 %id,
183 uploaded = stats.uploaded,
184 failed = stats.failed,
185 "cache drain completed"
186 );
187 }
188
189 let result = Message {
190 id,
191 cache_drain_result: Some(v1::CacheDrainResult {
192 error: protocol::error_or_empty(stats.last_error),
193 cache_queued: stats.pending,
194 cache_active: stats.active,
195 cache_uploaded: stats.uploaded,
196 cache_failed: stats.failed,
197 }),
198 ..Default::default()
199 };
200 let _ = out.send(result).await;
201}
202
203async fn run_poweroff(id: String, _req: v1::Poweroff, out: Sender<Message>) {
204 let result = Message {
205 id,
206 poweroff_result: Some(v1::PoweroffResult {
207 error: String::new(),
208 }),
209 ..Default::default()
210 };
211 let _ = out.send(result).await;
212
213 tokio::spawn(async move {
214 tokio::time::sleep(Duration::from_millis(100)).await;
215
216 // prefer a clean shutdown through the init system when one is around
217 // (systemd on NixOS, busybox/openrc elsewhere), then fall back to the
218 // raw reboot(2) syscall on minimal guests
219 for poweroff in [SYSTEMCTL_EXECUTABLE, "/sbin/poweroff", "/usr/sbin/poweroff"] {
220 if !std::path::Path::new(poweroff).exists() {
221 continue;
222 }
223 let mut spec = Spec::new(poweroff).timeout(Duration::from_secs(5));
224 if poweroff == SYSTEMCTL_EXECUTABLE {
225 spec = spec.args(["poweroff"]);
226 }
227 match command::run_capture(spec).await {
228 Ok(output) if output.success() => return,
229 Ok(output) => warn!(
230 %poweroff,
231 exit_code = output.exit.exit_code,
232 error = ?output.exit.error,
233 output = %output.combined_lossy(),
234 "poweroff command failed"
235 ),
236 Err(error) => warn!(%poweroff, %error, "poweroff command failed"),
237 }
238 }
239
240 // only ever returns on failure
241 let error =
242 nix::sys::reboot::reboot(nix::sys::reboot::RebootMode::RB_POWER_OFF).unwrap_err();
243 warn!(%error, "reboot(RB_POWER_OFF) syscall failed");
244 });
245}