alice.pds.demo.boltless.dev
This repository has no description
1use once_cell::sync::Lazy;
2use prost::Message as ProstMessage;
3use prost_reflect::DescriptorPool;
4use std::io;
5use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
6
7pub mod v1 {
8 include!("gen/spindle/agent/v1/spindle.agent.v1.rs");
9}
10
11pub use v1::Message;
12
13pub static DESCRIPTOR_POOL: Lazy<DescriptorPool> = Lazy::new(|| {
14 let bytes = include_bytes!("gen/file_descriptor_set.bin");
15 DescriptorPool::decode(&bytes[..]).unwrap()
16});
17
18macro_rules! impl_reflect {
19 ($($t:ident),* $(,)?) => {
20 $(
21 impl prost_reflect::ReflectMessage for v1::$t {
22 fn descriptor(&self) -> prost_reflect::MessageDescriptor {
23 DESCRIPTOR_POOL
24 .get_message_by_name(concat!("spindle.agent.v1.", stringify!($t)))
25 .unwrap()
26 }
27 }
28 )*
29 };
30}
31
32impl_reflect!(
33 Hello,
34 Init,
35 ExecStart,
36 ExecStderr,
37 ExecExit,
38 ActivateConfig,
39 ActivateConfigResult,
40 BuiltPaths,
41 CacheDrain,
42 CacheDrainResult,
43 Poweroff,
44 PoweroffResult,
45 Message,
46);
47
48pub const PROTOCOL_VERSION: u32 = 2;
49pub const DEFAULT_PORT: u32 = 10240;
50pub const MAX_MESSAGE_BYTES: usize = 1024 * 1024;
51
52#[macro_export]
53macro_rules! on_payload {
54 (ref $msg:expr, { $( $field:ident => $body:expr ),* $(,)? }) => {
55 #[allow(unused_variables)]
56 $(if let Some(ref $field) = $msg.$field { Some($body) } else)* { None }
57 };
58 ($msg:expr, { $( $field:ident => $body:expr ),* $(,)? }) => {
59 $(if let Some($field) = $msg.$field { Some($body) } else )* { None }
60 };
61}
62
63pub fn kind(msg: &Message) -> &'static str {
64 // todo(dawn): maybe eventually we should have a custom protoc plugin for
65 // generating an enum, right now not worth it, when we have more needs for
66 // it imo we can consider it again
67 on_payload!(ref msg, {
68 hello => "hello",
69 init => "init",
70 exec_start => "exec_start",
71 exec_stderr => "exec_stderr",
72 exec_exit => "exec_exit",
73 activate_config => "activate_config",
74 activate_config_result => "activate_config_result",
75 built_paths => "built_paths",
76 cache_drain => "cache_drain",
77 cache_drain_result => "cache_drain_result",
78 poweroff => "poweroff",
79 poweroff_result => "poweroff_result",
80 })
81 .unwrap_or_else(|| unreachable!("validated message has no payload"))
82}
83
84pub fn error_or_empty(error: Option<String>) -> String {
85 error.filter(|error| !error.is_empty()).unwrap_or_default()
86}
87
88pub async fn write_message<W: AsyncWrite + Unpin>(writer: &mut W, msg: &Message) -> io::Result<()> {
89 if let Err(err) = prost_protovalidate::validate(msg) {
90 return Err(io::Error::new(
91 io::ErrorKind::InvalidData,
92 format!("validate agent message: {err}"),
93 ));
94 }
95
96 let mut data = Vec::new();
97 msg.encode(&mut data)
98 .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?;
99 if data.len() > MAX_MESSAGE_BYTES {
100 return Err(io::Error::new(
101 io::ErrorKind::InvalidData,
102 format!("agent message exceeded {MAX_MESSAGE_BYTES} bytes"),
103 ));
104 }
105
106 writer.write_all(&(data.len() as u32).to_be_bytes()).await?;
107 writer.write_all(&data).await?;
108 writer.flush().await
109}
110
111pub async fn read_message<R: AsyncRead + Unpin>(reader: &mut R) -> io::Result<Option<Message>> {
112 let Some(header) = read_header(reader).await? else {
113 return Ok(None);
114 };
115 let size = u32::from_be_bytes(header) as usize;
116 if size > MAX_MESSAGE_BYTES {
117 return Err(io::Error::new(
118 io::ErrorKind::InvalidData,
119 format!("agent message exceeded {MAX_MESSAGE_BYTES} bytes"),
120 ));
121 }
122
123 let mut data = vec![0; size];
124 reader.read_exact(&mut data).await?;
125 let msg = Message::decode(&data[..])
126 .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?;
127
128 if let Err(err) = prost_protovalidate::validate(&msg) {
129 return Err(io::Error::new(
130 io::ErrorKind::InvalidData,
131 format!("validate agent message: {err}"),
132 ));
133 }
134
135 Ok(Some(msg))
136}
137
138async fn read_header<R: AsyncRead + Unpin>(reader: &mut R) -> io::Result<Option<[u8; 4]>> {
139 let mut header = [0; 4];
140 let mut read = 0;
141 while read < header.len() {
142 match reader.read(&mut header[read..]).await {
143 Ok(0) if read == 0 => return Ok(None),
144 Ok(0) => {
145 return Err(io::Error::new(
146 io::ErrorKind::UnexpectedEof,
147 "partial agent message header",
148 ));
149 }
150 Ok(n) => read += n,
151 Err(err) if err.kind() == io::ErrorKind::Interrupted => {}
152 Err(err) => return Err(err),
153 }
154 }
155 Ok(Some(header))
156}
157
158#[cfg(test)]
159mod tests {
160 use super::*;
161
162 #[tokio::test]
163 async fn round_trips_protobuf_message() {
164 let msg = Message {
165 id: "built-paths".to_owned(),
166 built_paths: Some(v1::BuiltPaths {
167 paths: vec!["/nix/store/abc-package".to_owned()],
168 reason: "post_build_hook".to_owned(),
169 }),
170 ..Default::default()
171 };
172
173 let mut encoded = Vec::new();
174 write_message(&mut encoded, &msg).await.unwrap();
175
176 let decoded = read_message(&mut &encoded[..]).await.unwrap().unwrap();
177 assert!(decoded.built_paths.is_some());
178 if let Some(p) = decoded.built_paths {
179 assert_eq!(p.paths, ["/nix/store/abc-package"]);
180 assert_eq!(p.reason, "post_build_hook");
181 }
182 }
183
184 #[test]
185 fn validates_messages() {
186 // 1. valid message (exactly one field set)
187 let valid = Message {
188 id: "test-1".to_owned(),
189 hello: Some(v1::Hello::default()),
190 ..Default::default()
191 };
192 assert!(prost_protovalidate::validate(&valid).is_ok());
193
194 // 2. invalid message (zero fields set)
195 let invalid_zero = Message {
196 id: "test-2".to_owned(),
197 ..Default::default()
198 };
199 assert!(prost_protovalidate::validate(&invalid_zero).is_err());
200
201 // 3. invalid message (multiple fields set)
202 let invalid_multi = Message {
203 id: "test-3".to_owned(),
204 hello: Some(v1::Hello::default()),
205 init: Some(v1::Init::default()),
206 ..Default::default()
207 };
208 assert!(prost_protovalidate::validate(&invalid_multi).is_err());
209 }
210}