Skip to main content

rpkl_jdx/api/
reader.rs

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    /// Scheme returns the scheme part of the URL that this reader can read.
14    /// The value should be the URI scheme up to (not including) `:`
15    fn scheme(&self) -> &str;
16
17    /// Tells whether the path part of ths URI has a
18    /// [hier-part](https://datatracker.ietf.org/doc/html/rfc3986#section-3).
19    ///
20    /// An example of a hierarchical URI is `file:///path/to/my/file`, where
21    /// `/path/to/my/file` designates a nested path through the `/` character.
22    ///
23    /// An example of a non-hierarchical URI is `pkl.base`, where the `base` does not denote
24    /// any form of hierarchy.
25    fn has_hierarchical_uris(&self) -> bool {
26        false
27    }
28
29    /// Tells whether this reader supports globbing.
30    fn is_globbable(&self) -> bool {
31        false
32    }
33
34    /// Read the contents of the resource at the given URI.
35    fn read(&self, uri: &str) -> Result<Vec<u8>, Box<dyn std::error::Error>>;
36
37    /// List the contents of the resource at the given URI.
38    fn list(&self, uri: &str) -> Result<Vec<PathElements>, Box<dyn std::error::Error>>;
39}
40
41pub trait PklModuleReader {
42    /// Scheme returns the scheme part of the URL that this reader can read.
43    /// The value should be the URI scheme up to (not including) `:`
44    fn scheme(&self) -> &str;
45
46    /// Tells whether the path part of ths URI has a
47    /// [hier-part](https://datatracker.ietf.org/doc/html/rfc3986#section-3).
48    ///
49    /// An example of a hierarchical URI is `file:///path/to/my/file`, where
50    /// `/path/to/my/file` designates a nested path through the `/` character.
51    ///
52    /// An example of a non-hierarchical URI is `pkl.base`, where the `base` does not denote
53    /// any form of hierarchy.
54    fn has_hierarchical_uris(&self) -> bool {
55        false
56    }
57
58    /// Tells whether this reader supports globbing.
59    fn is_globbable(&self) -> bool {
60        false
61    }
62
63    /// Tells whether the module is local to the system.
64    ///
65    /// A local resource that [hasHierarchicalUris] supports triple-dot imports.
66    fn is_local(&self) -> bool;
67
68    /// Read the contents of the module at the given URI.
69    fn read(&self, uri: &str) -> Result<String, Box<dyn std::error::Error>>;
70
71    /// List the contents of the module at the given URI.
72    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
132// TODO: there's a lot of duplicated code here
133// could be refactored, but the boilerplate needed and added complexity prob isn't worth it
134
135pub(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    // TODO: could add `with-serde` feature to rmpv to make this easier
143    // but might be overkill for messages with a small number of fields
144
145    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
404// Helper function to extract fields from response map
405fn 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}