Skip to main content

prns_runtime/runtime/
rnx.rs

1use prns_core::identity::IdentityHash;
2use prns_core::rnx::{
3    decode_execution_request_ref, encode_execution_result_into, EncodeExecutionResultError,
4    ExecutedCommandRef, ExecutionConclusion, ExecutionRequestRef, ExecutionResultRef,
5    RnxEncodeSink, MAX_RETURNED_STREAM_BYTES,
6};
7use prns_core::wire::DestinationHash;
8
9use super::request_endpoints::{
10    Decline, RequestContext, RequestEndpoint, RequestEndpointPolicy, ResponseCapacityExceeded,
11};
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum RnxAuthorization {
15    DenyAll,
16    AllowList(&'static [IdentityHash]),
17    Public,
18}
19
20impl RnxAuthorization {
21    const fn route_policy(self) -> RequestEndpointPolicy {
22        match self {
23            Self::DenyAll => RequestEndpointPolicy::AllowNone,
24            Self::AllowList(identities) => RequestEndpointPolicy::AllowList(identities),
25            Self::Public => RequestEndpointPolicy::AllowAll,
26        }
27    }
28}
29
30#[derive(Debug, Clone, Copy, PartialEq)]
31pub enum RnxCompletion {
32    NotExecuted {
33        started_at: f64,
34    },
35    Executed {
36        return_code: Option<i32>,
37        started_at: f64,
38        conclusion: ExecutionConclusion,
39    },
40}
41
42pub trait RnxOutputBuffer {
43    fn put(&mut self, bytes: &[u8]) -> usize;
44    fn as_slice(&self) -> &[u8];
45}
46
47impl<const N: usize> RnxOutputBuffer for heapless::Vec<u8, N> {
48    fn put(&mut self, bytes: &[u8]) -> usize {
49        let accepted = bytes.len().min(self.capacity().saturating_sub(self.len()));
50        let _ = self.extend_from_slice(&bytes[..accepted]);
51        accepted
52    }
53
54    fn as_slice(&self) -> &[u8] {
55        self.as_slice()
56    }
57}
58
59impl RnxOutputBuffer for alloc::vec::Vec<u8> {
60    fn put(&mut self, bytes: &[u8]) -> usize {
61        self.extend_from_slice(bytes);
62        bytes.len()
63    }
64
65    fn as_slice(&self) -> &[u8] {
66        self.as_slice()
67    }
68}
69
70pub trait RnxOutputStorage: Default {
71    fn buffers(&mut self) -> (&mut dyn RnxOutputBuffer, &mut dyn RnxOutputBuffer);
72}
73
74pub struct FixedRnxOutput<const STDOUT: usize, const STDERR: usize> {
75    stdout: heapless::Vec<u8, STDOUT>,
76    stderr: heapless::Vec<u8, STDERR>,
77}
78
79impl<const STDOUT: usize, const STDERR: usize> Default for FixedRnxOutput<STDOUT, STDERR> {
80    fn default() -> Self {
81        Self {
82            stdout: heapless::Vec::new(),
83            stderr: heapless::Vec::new(),
84        }
85    }
86}
87
88impl<const STDOUT: usize, const STDERR: usize> RnxOutputStorage for FixedRnxOutput<STDOUT, STDERR> {
89    fn buffers(&mut self) -> (&mut dyn RnxOutputBuffer, &mut dyn RnxOutputBuffer) {
90        (&mut self.stdout, &mut self.stderr)
91    }
92}
93
94#[derive(Default)]
95pub struct HeapRnxOutput {
96    stdout: alloc::vec::Vec<u8>,
97    stderr: alloc::vec::Vec<u8>,
98}
99
100impl RnxOutputStorage for HeapRnxOutput {
101    fn buffers(&mut self) -> (&mut dyn RnxOutputBuffer, &mut dyn RnxOutputBuffer) {
102        (&mut self.stdout, &mut self.stderr)
103    }
104}
105
106struct CapturedOutput<'a> {
107    buffer: &'a mut dyn RnxOutputBuffer,
108    returned_limit: Option<u64>,
109    total: u64,
110}
111
112impl CapturedOutput<'_> {
113    fn write(&mut self, bytes: &[u8]) {
114        self.total = self.total.saturating_add(bytes.len() as u64);
115        let remaining = self.returned_limit.map_or(u64::MAX, |limit| {
116            limit.saturating_sub(self.buffer.as_slice().len() as u64)
117        });
118        let accepted = usize::try_from(remaining)
119            .unwrap_or(usize::MAX)
120            .min(bytes.len());
121        self.buffer.put(&bytes[..accepted]);
122    }
123
124    fn observe_total(&mut self, total: u64) {
125        self.total = self.total.max(total);
126    }
127}
128
129pub struct RnxOutput<'a> {
130    stdout: CapturedOutput<'a>,
131    stderr: CapturedOutput<'a>,
132}
133
134impl<'a> RnxOutput<'a> {
135    pub fn new<T: RnxOutputStorage>(
136        storage: &'a mut T,
137        stdout_limit: Option<u64>,
138        stderr_limit: Option<u64>,
139    ) -> Self {
140        let (stdout, stderr) = storage.buffers();
141        let stdout_limit = Some(
142            stdout_limit
143                .unwrap_or(MAX_RETURNED_STREAM_BYTES as u64)
144                .min(MAX_RETURNED_STREAM_BYTES as u64),
145        );
146        let stderr_limit = Some(
147            stderr_limit
148                .unwrap_or(MAX_RETURNED_STREAM_BYTES as u64)
149                .min(MAX_RETURNED_STREAM_BYTES as u64),
150        );
151        Self {
152            stdout: CapturedOutput {
153                buffer: stdout,
154                returned_limit: stdout_limit,
155                total: 0,
156            },
157            stderr: CapturedOutput {
158                buffer: stderr,
159                returned_limit: stderr_limit,
160                total: 0,
161            },
162        }
163    }
164
165    pub fn stdout(&mut self, bytes: &[u8]) {
166        self.stdout.write(bytes);
167    }
168
169    pub fn stderr(&mut self, bytes: &[u8]) {
170        self.stderr.write(bytes);
171    }
172
173    pub fn observe_total_stdout(&mut self, total: u64) {
174        self.stdout.observe_total(total);
175    }
176
177    pub fn observe_total_stderr(&mut self, total: u64) {
178        self.stderr.observe_total(total);
179    }
180
181    #[must_use]
182    pub fn stdout_bytes(&self) -> &[u8] {
183        self.stdout.buffer.as_slice()
184    }
185
186    #[must_use]
187    pub fn stderr_bytes(&self) -> &[u8] {
188        self.stderr.buffer.as_slice()
189    }
190
191    #[must_use]
192    pub fn total_stdout(&self) -> u64 {
193        self.stdout.total
194    }
195
196    #[must_use]
197    pub fn total_stderr(&self) -> u64 {
198        self.stderr.total
199    }
200
201    fn result(&self, completion: RnxCompletion) -> ExecutionResultRef<'_> {
202        match completion {
203            RnxCompletion::NotExecuted { started_at } => {
204                ExecutionResultRef::NotExecuted { started_at }
205            }
206            RnxCompletion::Executed {
207                return_code,
208                started_at,
209                conclusion,
210            } => ExecutionResultRef::Executed(ExecutedCommandRef {
211                return_code,
212                stdout: self.stdout.buffer.as_slice(),
213                stderr: self.stderr.buffer.as_slice(),
214                total_stdout: self.stdout.total,
215                total_stderr: self.stderr.total,
216                started_at,
217                conclusion,
218            }),
219        }
220    }
221}
222
223#[allow(async_fn_in_trait)]
224pub trait RnxCommandHandler<State> {
225    const AUTHORIZATION: RnxAuthorization = RnxAuthorization::DenyAll;
226    type Output: RnxOutputStorage;
227
228    fn destination(state: &State) -> DestinationHash;
229
230    async fn execute(
231        state: &State,
232        request: ExecutionRequestRef<'_>,
233        output: &mut RnxOutput<'_>,
234    ) -> RnxCompletion;
235}
236
237impl<State, Endpoint> RequestEndpoint<State> for Endpoint
238where
239    Endpoint: RnxCommandHandler<State>,
240{
241    const ENDPOINT_ID: &'static str = prns_core::rnx::COMMAND_PATH;
242    const POLICY: RequestEndpointPolicy = Endpoint::AUTHORIZATION.route_policy();
243
244    async fn handle(mut context: RequestContext<'_, State>) -> Result<(), Decline> {
245        if context.destination != Endpoint::destination(context.state) {
246            return Err(Decline::Ignore);
247        }
248        let request = decode_execution_request_ref(context.data).map_err(|_| Decline::Ignore)?;
249        let mut storage = Endpoint::Output::default();
250        let mut output = RnxOutput::new(&mut storage, request.stdout_limit, request.stderr_limit);
251        let completion = Endpoint::execute(context.state, request, &mut output).await;
252        let result = output.result(completion);
253        encode_execution_result_into(result, &mut ContextSink(&mut context)).map_err(|error| {
254            match error {
255                EncodeExecutionResultError::Codec(_) => Decline::Ignore,
256                EncodeExecutionResultError::Sink(_) => Decline::ResponseTooLarge,
257            }
258        })
259    }
260}
261
262struct ContextSink<'a, 'request, State>(&'a mut RequestContext<'request, State>);
263
264impl<State> RnxEncodeSink for ContextSink<'_, '_, State> {
265    type Error = ResponseCapacityExceeded;
266
267    fn put(&mut self, bytes: &[u8]) -> Result<(), Self::Error> {
268        self.0.write_packed(bytes).map(|_| ())
269    }
270}
271
272#[cfg(test)]
273mod tests {
274    use super::*;
275    use prns_core::engine::InstantMillis;
276    use prns_core::identity::IdentityHash;
277    use prns_core::rnx::{
278        decode_execution_result, encode_execution_request, ExecutionRequest, ExecutionResult,
279    };
280    use prns_core::routing::links::request::RequestId;
281    use prns_core::routing::links::LinkId;
282    use prns_core::routing::request_handlers::RequestPathHash;
283    use prns_core::units::RttMillis;
284
285    const DESTINATION: DestinationHash = DestinationHash::new([0x44; 16]);
286    const ADMIN: IdentityHash = IdentityHash::new([0x55; 16]);
287
288    struct App;
289    struct DeniedEndpoint;
290    struct RnxEndpoint;
291
292    impl RnxCommandHandler<App> for DeniedEndpoint {
293        type Output = FixedRnxOutput<0, 0>;
294
295        fn destination(_state: &App) -> DestinationHash {
296            DESTINATION
297        }
298
299        async fn execute(
300            _state: &App,
301            _request: ExecutionRequestRef<'_>,
302            _output: &mut RnxOutput<'_>,
303        ) -> RnxCompletion {
304            RnxCompletion::NotExecuted { started_at: 1.0 }
305        }
306    }
307
308    impl RnxCommandHandler<App> for RnxEndpoint {
309        const AUTHORIZATION: RnxAuthorization = RnxAuthorization::AllowList(&[ADMIN]);
310        type Output = FixedRnxOutput<4, 2>;
311
312        fn destination(_state: &App) -> DestinationHash {
313            DESTINATION
314        }
315
316        async fn execute(
317            _state: &App,
318            request: ExecutionRequestRef<'_>,
319            output: &mut RnxOutput<'_>,
320        ) -> RnxCompletion {
321            if request.command != "status" {
322                return RnxCompletion::NotExecuted { started_at: 1.0 };
323            }
324            output.stdout(b"ready");
325            output.stderr(b"warn");
326            RnxCompletion::Executed {
327                return_code: Some(0),
328                started_at: 1.0,
329                conclusion: ExecutionConclusion::CompletedAt(2.0),
330            }
331        }
332    }
333
334    #[test]
335    fn the_endpoint_type_is_the_route_and_bounds_its_output() {
336        futures_executor::block_on(async {
337            async fn dispatch<R: super::super::request_endpoints::RequestEndpointSet<App>>(
338                _endpoints: &R,
339                destination: DestinationHash,
340                sink: &mut dyn super::super::request_endpoints::ResponseSink,
341            ) -> Result<(), Decline> {
342                let request = ExecutionRequest {
343                    command: alloc::string::String::from("status"),
344                    timeout_seconds: None,
345                    stdout_limit: None,
346                    stderr_limit: None,
347                    stdin: None,
348                };
349                let data = encode_execution_request(&request).unwrap();
350                let inbound = super::super::request_endpoints::InboundRequest::new(
351                    destination,
352                    LinkId::new([1; 16]),
353                    RequestId([2; 16]),
354                    Some(ADMIN),
355                    InstantMillis(3),
356                    RttMillis::new(4),
357                    &data,
358                );
359                super::super::request_endpoints::dispatch_request::<App, R>(
360                    &App,
361                    RequestPathHash::of(prns_core::rnx::COMMAND_PATH),
362                    inbound,
363                    sink,
364                )
365                .await
366            }
367
368            let endpoints = crate::request_endpoints![RnxEndpoint];
369            assert_eq!(DeniedEndpoint::POLICY, RequestEndpointPolicy::AllowNone);
370            assert_eq!(
371                RnxEndpoint::POLICY,
372                RequestEndpointPolicy::AllowList(&[ADMIN])
373            );
374            let mut encoded = heapless::Vec::<u8, 128>::new();
375            assert_eq!(
376                dispatch(&endpoints, DESTINATION, &mut encoded).await,
377                Ok(())
378            );
379            let ExecutionResult::Executed(result) =
380                decode_execution_result(encoded.as_slice()).unwrap()
381            else {
382                panic!("executed result");
383            };
384            assert_eq!(result.stdout, b"read");
385            assert_eq!(result.stderr, b"wa");
386            assert_eq!(result.total_stdout, 5);
387            assert_eq!(result.total_stderr, 4);
388
389            let mut wrong_destination = heapless::Vec::<u8, 128>::new();
390            assert_eq!(
391                dispatch(
392                    &endpoints,
393                    DestinationHash::new([0x66; 16]),
394                    &mut wrong_destination,
395                )
396                .await,
397                Err(Decline::Ignore)
398            );
399            assert!(wrong_destination.is_empty());
400        });
401    }
402}