1pub use crate::internal::msgapi::PathElements;
2use crate::internal::msgapi::{
3 incoming::PklServerMessage,
4 outgoing::{
5 ListModulesResponse, ListResourcesResponse, ReadModuleResponse, ReadResourceResponse,
6 },
7 PklMessage,
8};
9use crate::utils::macros::_warn;
10use std::io::Write;
11
12pub trait PklResourceReader {
13 fn scheme(&self) -> &str;
16
17 fn has_hierarchical_uris(&self) -> bool {
26 false
27 }
28
29 fn is_globbable(&self) -> bool {
31 false
32 }
33
34 fn read(&self, uri: &str) -> Result<Vec<u8>, Box<dyn std::error::Error>>;
36
37 fn list(&self, uri: &str) -> Result<Vec<PathElements>, Box<dyn std::error::Error>>;
39}
40
41pub trait PklModuleReader {
42 fn scheme(&self) -> &str;
45
46 fn has_hierarchical_uris(&self) -> bool {
55 false
56 }
57
58 fn is_globbable(&self) -> bool {
60 false
61 }
62
63 fn is_local(&self) -> bool;
67
68 fn read(&self, uri: &str) -> Result<String, Box<dyn std::error::Error>>;
70
71 fn list(&self, uri: &str) -> Result<Vec<PathElements>, Box<dyn std::error::Error>>;
73}
74
75pub trait IntoResourceReaders {
76 fn into_readers(self) -> Vec<Box<dyn PklResourceReader>>;
77}
78
79pub trait IntoModuleReaders {
80 fn into_readers(self) -> Vec<Box<dyn PklModuleReader>>;
81}
82
83macro_rules! impl_into_readers {
84 (resource, $(($type:ident)),+) => {
85 #[allow(non_snake_case)]
86 impl<$($type),+> IntoResourceReaders for ($($type),+)
87 where
88 $($type: PklResourceReader + 'static),+
89 {
90 fn into_readers(self) -> Vec<Box<dyn PklResourceReader>> {
91 let ($($type),+) = self;
92 vec![$(Box::new($type)),+]
93 }
94 }
95 };
96 (module, $(($type:ident)),+) => {
97 #[allow(non_snake_case)]
98 impl<$($type),+> IntoModuleReaders for ($($type),+)
99 where
100 $($type: PklModuleReader + 'static),+
101 {
102 fn into_readers(self) -> Vec<Box<dyn PklModuleReader>> {
103 let ($($type),+) = self;
104 vec![$(Box::new($type)),+]
105 }
106 }
107 };
108}
109
110impl<T: PklResourceReader + 'static> IntoResourceReaders for T {
111 fn into_readers(self) -> Vec<Box<dyn PklResourceReader>> {
112 vec![Box::new(self)]
113 }
114}
115
116impl<T: PklModuleReader + 'static> IntoModuleReaders for T {
117 fn into_readers(self) -> Vec<Box<dyn PklModuleReader>> {
118 vec![Box::new(self)]
119 }
120}
121
122impl_into_readers!(resource, (T1), (T2));
123impl_into_readers!(resource, (T1), (T2), (T3));
124impl_into_readers!(resource, (T1), (T2), (T3), (T4));
125impl_into_readers!(resource, (T1), (T2), (T3), (T4), (T5));
126
127impl_into_readers!(module, (T1), (T2));
128impl_into_readers!(module, (T1), (T2), (T3));
129impl_into_readers!(module, (T1), (T2), (T3), (T4));
130impl_into_readers!(module, (T1), (T2), (T3), (T4), (T5));
131
132pub(crate) fn handle_list_resources<W: Write>(
136 resource_readers: &[Box<dyn PklResourceReader>],
137 msg: &PklServerMessage,
138 writer: &mut W,
139) -> Result<(), Box<dyn std::error::Error>> {
140 let response = msg.response.as_map().unwrap();
141
142 let evaluator_id: i64 = extract_field(response, "evaluatorId")?;
146 let request_id: i64 = extract_field(response, "requestId")?;
147 let uri: &str = extract_field(response, "uri")?;
148
149 let uri_scheme = parse_scheme(uri).expect("Invalid URI, this is a bug");
150
151 let Some(reader) = resource_readers.iter().find(|r| r.scheme() == uri_scheme) else {
152 _warn!("No reader found for scheme: {:?}", uri);
153 writer.write_all(
154 &ListResourcesResponse {
155 request_id,
156 evaluator_id,
157 path_elements: None,
158 error: Some(format!("No reader found for scheme: {:?}", uri)),
159 }
160 .encode_msg()?,
161 )?;
162 writer.flush()?;
163 return Ok(());
164 };
165
166 let data = reader.list(uri);
167
168 let out_msg = match data {
169 Ok(elements) => ListResourcesResponse {
170 request_id,
171 evaluator_id,
172 path_elements: Some(elements),
173 error: None,
174 },
175 Err(e) => ListResourcesResponse {
176 request_id,
177 evaluator_id,
178 path_elements: None,
179 error: Some(e.to_string()),
180 },
181 };
182
183 writer.write_all(&out_msg.encode_msg()?)?;
184 writer.flush()?;
185
186 Ok(())
187}
188
189pub(crate) fn handle_list_modules<W: Write>(
190 module_readers: &[Box<dyn PklModuleReader>],
191 msg: &PklServerMessage,
192 writer: &mut W,
193) -> Result<(), Box<dyn std::error::Error>> {
194 let response = msg.response.as_map().unwrap();
195
196 let evaluator_id: i64 = extract_field(response, "evaluatorId")?;
197 let request_id: i64 = extract_field(response, "requestId")?;
198 let uri: &str = extract_field(response, "uri")?;
199
200 let uri_scheme = parse_scheme(uri).expect("Invalid URI, this is a bug");
201
202 let Some(reader) = module_readers.iter().find(|r| r.scheme() == uri_scheme) else {
203 _warn!("No reader found for scheme: {:?}", uri);
204 writer.write_all(
205 &ListModulesResponse {
206 request_id,
207 evaluator_id,
208 path_elements: None,
209 error: Some(format!("No reader found for scheme: {:?}", uri)),
210 }
211 .encode_msg()?,
212 )?;
213 writer.flush()?;
214 return Ok(());
215 };
216
217 let data = reader.list(uri);
218
219 let out_msg = match data {
220 Ok(elements) => ListModulesResponse {
221 request_id,
222 evaluator_id,
223 path_elements: Some(elements),
224 error: None,
225 },
226 Err(e) => ListModulesResponse {
227 request_id,
228 evaluator_id,
229 path_elements: None,
230 error: Some(e.to_string()),
231 },
232 };
233
234 writer.write_all(&out_msg.encode_msg()?)?;
235 writer.flush()?;
236
237 Ok(())
238}
239
240pub(crate) fn handle_read_resource<W: Write>(
241 resource_readers: &[Box<dyn PklResourceReader>],
242 msg: &PklServerMessage,
243 writer: &mut W,
244) -> Result<(), Box<dyn std::error::Error>> {
245 let response = msg.response.as_map().unwrap();
246
247 let evaluator_id: i64 = extract_field(response, "evaluatorId")?;
248 let request_id: i64 = extract_field(response, "requestId")?;
249 let uri: &str = extract_field(response, "uri")?;
250
251 let uri_scheme = parse_scheme(uri).expect("Invalid URI, this is a bug");
252
253 let Some(reader) = resource_readers.iter().find(|r| r.scheme() == uri_scheme) else {
254 _warn!("No reader found for scheme: {:?}", uri);
255 writer.write_all(
256 &ReadResourceResponse {
257 request_id,
258 evaluator_id,
259 contents: None,
260 error: Some(format!("No reader found for scheme: {:?}", uri)),
261 }
262 .encode_msg()?,
263 )?;
264 writer.flush()?;
265 return Ok(());
266 };
267
268 let data = reader.read(uri);
269
270 let out_msg = match data {
271 Ok(data) => ReadResourceResponse {
272 request_id,
273 evaluator_id,
274 contents: Some(data),
275 error: None,
276 },
277 Err(e) => ReadResourceResponse {
278 request_id,
279 evaluator_id,
280 contents: None,
281 error: Some(e.to_string()),
282 },
283 };
284
285 let serialized = out_msg.encode_msg()?;
286
287 writer.write_all(&serialized)?;
288 writer.flush()?;
289
290 Ok(())
291}
292
293pub(crate) fn handle_read_module<W: Write>(
294 module_readers: &[Box<dyn PklModuleReader>],
295 msg: &PklServerMessage,
296 writer: &mut W,
297) -> Result<(), Box<dyn std::error::Error>> {
298 let response = msg.response.as_map().unwrap();
299
300 let evaluator_id: i64 = extract_field(response, "evaluatorId")?;
301 let request_id: i64 = extract_field(response, "requestId")?;
302 let uri: &str = extract_field(response, "uri")?;
303
304 let uri_scheme = parse_scheme(uri).expect("Invalid URI, this is a bug");
305
306 let Some(reader) = module_readers.iter().find(|r| r.scheme() == uri_scheme) else {
307 _warn!("No reader found for scheme: {:?}", uri);
308 writer.write_all(
309 &ReadModuleResponse {
310 request_id,
311 evaluator_id,
312 contents: None,
313 error: Some(format!("No reader found for scheme: {:?}", uri)),
314 }
315 .encode_msg()?,
316 )?;
317 writer.flush()?;
318 return Ok(());
319 };
320
321 let data = reader.read(uri);
322
323 let out_msg = match data {
324 Ok(data) => ReadModuleResponse {
325 request_id,
326 evaluator_id,
327 contents: Some(data),
328 error: None,
329 },
330 Err(e) => ReadModuleResponse {
331 request_id,
332 evaluator_id,
333 contents: None,
334 error: Some(e.to_string()),
335 },
336 };
337
338 let serialized = out_msg.encode_msg()?;
339
340 writer.write_all(&serialized)?;
341 writer.flush()?;
342
343 Ok(())
344}
345
346struct MapValue<'a>(&'a rmpv::Value);
347
348impl<'a> TryFrom<MapValue<'a>> for i64 {
349 type Error = Box<dyn std::error::Error>;
350
351 fn try_from(value: MapValue<'a>) -> Result<Self, Self::Error> {
352 match value.0 {
353 rmpv::Value::Integer(n) => n.as_i64().ok_or_else(|| "Failed to convert to i64".into()),
354 _ => Err("Expected integer value".into()),
355 }
356 }
357}
358
359impl<'a> TryFrom<MapValue<'a>> for &'a str {
360 type Error = Box<dyn std::error::Error>;
361
362 fn try_from(value: MapValue<'a>) -> Result<Self, Self::Error> {
363 match value.0 {
364 rmpv::Value::String(s) => s
365 .as_str()
366 .ok_or_else(|| "Failed to get str from string".into()),
367 _ => Err("Expected string value".into()),
368 }
369 }
370}
371
372impl<'a> TryFrom<MapValue<'a>> for String {
373 type Error = Box<dyn std::error::Error>;
374
375 fn try_from(value: MapValue<'a>) -> Result<Self, Self::Error> {
376 match value.0 {
377 rmpv::Value::String(s) => Ok(s
378 .as_str()
379 .ok_or_else(|| "Failed to get str from string")?
380 .to_owned()),
381 _ => Err("Expected string value".into()),
382 }
383 }
384}
385
386fn parse_scheme(uri: &str) -> Option<&str> {
387 match uri.find(':') {
388 Some(pos) => {
389 let scheme = &uri[..pos];
390 if !scheme.is_empty()
391 && scheme
392 .chars()
393 .all(|c| c.is_ascii_alphanumeric() || c == '+' || c == '.' || c == '-')
394 {
395 Some(scheme)
396 } else {
397 None
398 }
399 }
400 None => None,
401 }
402}
403
404fn extract_field<'a, T>(
406 map: &'a [(rmpv::Value, rmpv::Value)],
407 field: &str,
408) -> Result<T, Box<dyn std::error::Error>>
409where
410 T: TryFrom<MapValue<'a>, Error = Box<dyn std::error::Error>>,
411{
412 map.iter()
413 .find(|(k, _)| k.as_str() == Some(field))
414 .map(|(_, v)| MapValue(v))
415 .ok_or_else(|| format!("Field not found in message: {}", field).into())
416 .and_then(|v| v.try_into())
417}