capsule_runtime/guest_worker/
mod.rs1#[cfg(all(feature = "escape-probe", target_os = "linux"))]
21#[allow(unsafe_code)]
22mod escape_probe;
23mod http;
24mod launch;
25mod remote;
26mod supervisor;
27pub(crate) mod wire;
28mod worker;
29
30use std::ffi::OsString;
31use std::sync::atomic::{AtomicU8, Ordering};
32
33use serde::{Deserialize, Serialize};
34
35use crate::errors::RuntimeError;
36
37pub(crate) use http::RemoteHttp;
38pub(crate) use launch::{Launcher, WorkerSandbox};
39pub(crate) use remote::RemoteProcessDriver;
40pub use supervisor::GuestWorkers;
41pub(crate) use wire::StreamEvent;
42pub use worker::RemoteLink;
43
44pub const PROTOCOL_VERSION: u32 = 1;
46
47pub const ENTRY: &str = "__guest-worker";
49
50#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
53pub enum WorkerKind {
54 Driver,
57 Tools,
59 Hooks,
61}
62
63impl WorkerKind {
64 pub fn as_str(self) -> &'static str {
65 match self {
66 Self::Driver => "driver",
67 Self::Tools => "tools",
68 Self::Hooks => "hooks",
69 }
70 }
71
72 pub fn parse(name: &str) -> Option<Self> {
73 match name {
74 "driver" => Some(Self::Driver),
75 "tools" => Some(Self::Tools),
76 "hooks" => Some(Self::Hooks),
77 _ => None,
78 }
79 }
80
81 pub fn for_runtime(runtime: murmur_artifact::ArtifactRuntime) -> Option<Self> {
84 use murmur_artifact::ArtifactRuntime;
85 match runtime {
86 ArtifactRuntime::Tool => Some(Self::Tools),
87 ArtifactRuntime::Driver => Some(Self::Driver),
88 ArtifactRuntime::Hook => Some(Self::Hooks),
89 ArtifactRuntime::Skill => None,
90 }
91 }
92
93 fn runtime_threads(self) -> usize {
95 match self {
96 Self::Tools => 4,
97 Self::Driver | Self::Hooks => 2,
98 }
99 }
100}
101
102const PLACEMENT_UNSPLIT: u8 = 0;
103const PLACEMENT_SUPERVISOR: u8 = 1;
104const PLACEMENT_WORKER: u8 = 2;
105
106static PLACEMENT: AtomicU8 = AtomicU8::new(PLACEMENT_UNSPLIT);
108
109#[derive(Debug, Clone, Copy, PartialEq, Eq)]
111pub(crate) enum Placement {
112 Unsplit,
114 Supervisor,
116 Worker,
118}
119
120pub fn use_worker_processes() {
123 let _ = PLACEMENT.compare_exchange(
124 PLACEMENT_UNSPLIT,
125 PLACEMENT_SUPERVISOR,
126 Ordering::SeqCst,
127 Ordering::SeqCst,
128 );
129}
130
131#[cfg(test)]
132thread_local! {
133 static TEST_PLACEMENT: std::cell::Cell<Option<Placement>> = const { std::cell::Cell::new(None) };
135}
136
137#[cfg(test)]
140pub(crate) fn with_placement<R>(placement: Placement, f: impl FnOnce() -> R) -> R {
141 TEST_PLACEMENT.with(|cell| cell.set(Some(placement)));
142 let result = f();
143 TEST_PLACEMENT.with(|cell| cell.set(None));
144 result
145}
146
147pub(crate) fn current_placement() -> Placement {
148 #[cfg(test)]
149 if let Some(placement) = TEST_PLACEMENT.with(std::cell::Cell::get) {
150 return placement;
151 }
152 match PLACEMENT.load(Ordering::SeqCst) {
153 PLACEMENT_SUPERVISOR => Placement::Supervisor,
154 PLACEMENT_WORKER => Placement::Worker,
155 _ => Placement::Unsplit,
156 }
157}
158
159pub(crate) fn session_launcher() -> Option<Launcher> {
161 match current_placement() {
162 Placement::Supervisor => Some(Launcher::own_image()),
163 Placement::Unsplit | Placement::Worker => None,
164 }
165}
166
167#[derive(Debug, Clone, Copy, PartialEq, Eq)]
169pub(crate) enum StoreSite {
170 Tool,
171 Driver,
172 Hook,
173 ProcessDriver,
174 Capsule,
175}
176
177impl StoreSite {
178 fn describe(self) -> &'static str {
179 match self {
180 Self::Tool => "a tool",
181 Self::Driver => "an inference driver",
182 Self::Hook => "a hook",
183 Self::ProcessDriver => "a process driver",
184 Self::Capsule => "a script capsule's component",
185 }
186 }
187}
188
189pub(crate) fn store_allowed(placement: Placement, site: StoreSite) -> Result<(), RuntimeError> {
192 match placement {
193 Placement::Unsplit | Placement::Worker => Ok(()),
194 Placement::Supervisor => Err(RuntimeError::Runtime(format!(
195 "refused to create a store for {} in the supervisor: guests run only in guest workers",
196 site.describe()
197 ))),
198 }
199}
200
201pub(crate) fn check_store_placement(site: StoreSite) -> Result<(), RuntimeError> {
203 store_allowed(current_placement(), site)
204}
205
206#[derive(Debug, Clone, PartialEq, Eq)]
208pub enum WorkerEntry {
209 Kind(WorkerKind),
210 Misuse,
212}
213
214pub fn entry_kind(args: impl IntoIterator<Item = OsString>) -> Option<WorkerEntry> {
216 let mut args = args.into_iter().skip(1);
217 if args.next()?.to_str() != Some(ENTRY) {
218 return None;
219 }
220 let kind = args.next();
221 if args.next().is_some() {
222 return Some(WorkerEntry::Misuse);
223 }
224 Some(
225 kind.and_then(|kind| kind.to_str().and_then(WorkerKind::parse))
226 .map_or(WorkerEntry::Misuse, WorkerEntry::Kind),
227 )
228}
229
230pub fn worker_main(entry: WorkerEntry) -> i32 {
233 let kind = match entry {
234 WorkerEntry::Kind(kind) => kind,
235 WorkerEntry::Misuse => {
236 crate::diagnostic::to_stderr(&format!(
237 "mur {ENTRY} is an internal entry started by mur run; it is not a command"
238 ));
239 return 2;
240 }
241 };
242 let count = std::env::vars_os().count();
243 if count != 0 {
244 crate::diagnostic::to_stderr(&format!(
245 "mur {ENTRY} {kind}: refused to start: its environment holds {count} variable(s); a \
246 guest worker starts with an empty environment",
247 kind = kind.as_str()
248 ));
249 return 2;
250 }
251 let Some(stream) = connection_on_fd3() else {
252 crate::diagnostic::to_stderr(&format!(
253 "mur {ENTRY} {kind}: refused to start: fd 3 is not its connection; it is an internal \
254 entry started by mur run",
255 kind = kind.as_str()
256 ));
257 return 2;
258 };
259 PLACEMENT.store(PLACEMENT_WORKER, Ordering::SeqCst);
260 worker::serve(kind, stream)
261}
262
263#[allow(unsafe_code)]
264fn connection_on_fd3() -> Option<std::os::unix::net::UnixStream> {
265 use std::os::fd::FromRawFd;
266 let is_socket = unsafe {
268 let mut stat: libc::stat = std::mem::zeroed();
269 libc::fstat(3, &mut stat) == 0 && (stat.st_mode & libc::S_IFMT) == libc::S_IFSOCK
270 };
271 if !is_socket {
272 return None;
273 }
274 let stream = unsafe { std::os::unix::net::UnixStream::from_raw_fd(3) };
277 unsafe {
280 let flags = libc::fcntl(3, libc::F_GETFD);
281 if flags >= 0 {
282 libc::fcntl(3, libc::F_SETFD, flags | libc::FD_CLOEXEC);
283 }
284 }
285 Some(stream)
286}
287
288#[cfg(test)]
289mod tests;