sim_lib_server/transport/
backends.rs

1use std::{
2    collections::BTreeMap,
3    sync::{Arc, Mutex, OnceLock},
4};
5
6use sim_kernel::{Cx, Error, Result};
7
8#[cfg(feature = "wasm")]
9use sim_wasm_abi::{Frame as WasmFrame, WasmFrameLimits};
10
11#[cfg(feature = "wasm")]
12use crate::wasm::lookup_wasm_region;
13use crate::{EvalSite, ServerAddress, ServerFrame};
14
15use super::framing::endpoint_key;
16use super::{
17    ConnectionTransport, ServerTransport, decode_transport_frame, encode_transport_frame,
18    route_frame_bytes,
19};
20
21#[derive(Clone)]
22pub(crate) struct TransportEndpoint {
23    pub(crate) address: ServerAddress,
24    pub(crate) site: Arc<dyn EvalSite>,
25}
26
27#[derive(Default)]
28struct EndpointRegistry {
29    endpoints: BTreeMap<String, TransportEndpoint>,
30}
31
32fn endpoint_registry() -> &'static Mutex<EndpointRegistry> {
33    static REGISTRY: OnceLock<Mutex<EndpointRegistry>> = OnceLock::new();
34    REGISTRY.get_or_init(|| Mutex::new(EndpointRegistry::default()))
35}
36
37pub(crate) fn register_endpoint(endpoint: TransportEndpoint) -> Result<()> {
38    let mut registry = endpoint_registry()
39        .lock()
40        .map_err(|_| Error::HostError("endpoint registry mutex poisoned".to_owned()))?;
41    registry
42        .endpoints
43        .insert(endpoint_key(&endpoint.address), endpoint);
44    Ok(())
45}
46
47pub(crate) fn has_registered_endpoint(address: &ServerAddress) -> Result<bool> {
48    let registry = endpoint_registry()
49        .lock()
50        .map_err(|_| Error::HostError("endpoint registry mutex poisoned".to_owned()))?;
51    Ok(registry.endpoints.contains_key(&endpoint_key(address)))
52}
53
54pub(crate) fn unregister_endpoint(address: &ServerAddress) -> Result<()> {
55    let mut registry = endpoint_registry()
56        .lock()
57        .map_err(|_| Error::HostError("endpoint registry mutex poisoned".to_owned()))?;
58    registry.endpoints.remove(&endpoint_key(address));
59    Ok(())
60}
61
62pub(crate) fn lookup_endpoint(address: &ServerAddress) -> Result<TransportEndpoint> {
63    let registry = endpoint_registry()
64        .lock()
65        .map_err(|_| Error::HostError("endpoint registry mutex poisoned".to_owned()))?;
66    registry
67        .endpoints
68        .get(&endpoint_key(address))
69        .cloned()
70        .ok_or_else(|| {
71            Error::Eval(format!(
72                "no endpoint registered for {}",
73                address.kind_symbol()
74            ))
75        })
76}
77
78/// RAII guard for a registered loopback endpoint.
79///
80/// Dropping it, or calling [`close`](Self::close), unregisters the endpoint.
81pub struct LoopbackTransportEndpoint {
82    address: ServerAddress,
83}
84
85impl LoopbackTransportEndpoint {
86    /// Returns the registered endpoint address.
87    pub fn address(&self) -> &ServerAddress {
88        &self.address
89    }
90
91    /// Unregisters the endpoint explicitly.
92    pub fn close(&self) -> Result<()> {
93        unregister_endpoint(&self.address)
94    }
95}
96
97impl Drop for LoopbackTransportEndpoint {
98    fn drop(&mut self) {
99        let _ = unregister_endpoint(&self.address);
100    }
101}
102
103pub(crate) fn register_loopback_endpoint(
104    address: ServerAddress,
105    site: Arc<dyn EvalSite>,
106) -> Result<LoopbackTransportEndpoint> {
107    register_endpoint(TransportEndpoint {
108        address: address.clone(),
109        site,
110    })?;
111    Ok(LoopbackTransportEndpoint { address })
112}
113
114#[derive(Clone)]
115/// In-process transport that routes frames directly to a local eval site.
116///
117/// Sending a frame evaluates it against the site and buffers the reply for the
118/// next receive.
119pub struct LocalTransport {
120    address: ServerAddress,
121    site: Arc<dyn EvalSite>,
122    pending: Arc<Mutex<Option<ServerFrame>>>,
123}
124
125impl LocalTransport {
126    /// Creates a local transport for `address` backed by `site`.
127    pub fn new(address: ServerAddress, site: Arc<dyn EvalSite>) -> Self {
128        Self {
129            address,
130            site,
131            pending: Arc::new(Mutex::new(None)),
132        }
133    }
134}
135
136impl ServerTransport for LocalTransport {
137    fn address(&self) -> &ServerAddress {
138        &self.address
139    }
140
141    fn accept(&self, _cx: &mut Cx) -> Result<Box<dyn ConnectionTransport>> {
142        Ok(Box::new(self.clone()))
143    }
144
145    fn shutdown(&self, _cx: &mut Cx) -> Result<()> {
146        let mut pending = self
147            .pending
148            .lock()
149            .map_err(|_| Error::HostError("local transport mutex poisoned".to_owned()))?;
150        *pending = None;
151        Ok(())
152    }
153
154    fn accept_timeout(
155        &self,
156        _cx: &mut Cx,
157        _timeout: std::time::Duration,
158    ) -> Result<Option<Box<dyn ConnectionTransport>>> {
159        Ok(None)
160    }
161}
162
163impl ConnectionTransport for LocalTransport {
164    fn send_frame(&mut self, cx: &mut Cx, frame: ServerFrame) -> Result<()> {
165        let reply = route_frame_bytes(cx, &self.site, &encode_transport_frame(&frame)?)?;
166        let reply = decode_transport_frame(&reply)?;
167        let mut pending = self
168            .pending
169            .lock()
170            .map_err(|_| Error::HostError("local transport mutex poisoned".to_owned()))?;
171        *pending = Some(reply);
172        Ok(())
173    }
174
175    fn recv_frame(
176        &mut self,
177        _cx: &mut Cx,
178        _timeout: Option<std::time::Duration>,
179    ) -> Result<Option<ServerFrame>> {
180        let mut pending = self
181            .pending
182            .lock()
183            .map_err(|_| Error::HostError("local transport mutex poisoned".to_owned()))?;
184        Ok(pending.take())
185    }
186
187    fn close(&mut self, _cx: &mut Cx) -> Result<()> {
188        let mut pending = self
189            .pending
190            .lock()
191            .map_err(|_| Error::HostError("local transport mutex poisoned".to_owned()))?;
192        *pending = None;
193        Ok(())
194    }
195
196    fn as_any(&self) -> &dyn std::any::Any {
197        self
198    }
199}
200
201#[derive(Clone)]
202pub struct RegistryTransport {
203    address: ServerAddress,
204    pending: Arc<Mutex<Option<ServerFrame>>>,
205}
206
207impl RegistryTransport {
208    pub fn new(address: ServerAddress) -> Self {
209        Self {
210            address,
211            pending: Arc::new(Mutex::new(None)),
212        }
213    }
214}
215
216impl ServerTransport for RegistryTransport {
217    fn address(&self) -> &ServerAddress {
218        &self.address
219    }
220
221    fn accept(&self, _cx: &mut Cx) -> Result<Box<dyn ConnectionTransport>> {
222        let _ = lookup_endpoint(self.address())?;
223        Ok(Box::new(Self::new(self.address.clone())))
224    }
225
226    fn shutdown(&self, _cx: &mut Cx) -> Result<()> {
227        unregister_endpoint(&self.address)
228    }
229
230    fn accept_timeout(
231        &self,
232        _cx: &mut Cx,
233        _timeout: std::time::Duration,
234    ) -> Result<Option<Box<dyn ConnectionTransport>>> {
235        Ok(None)
236    }
237}
238
239impl ConnectionTransport for RegistryTransport {
240    fn send_frame(&mut self, cx: &mut Cx, frame: ServerFrame) -> Result<()> {
241        let endpoint = lookup_endpoint(&self.address)?;
242        let bytes = encode_transport_frame(&frame)?;
243        let reply = route_frame_bytes(cx, &endpoint.site, &bytes)?;
244        let reply = decode_transport_frame(&reply)?;
245        let mut pending = self
246            .pending
247            .lock()
248            .map_err(|_| Error::HostError("registry transport mutex poisoned".to_owned()))?;
249        *pending = Some(reply);
250        Ok(())
251    }
252
253    fn recv_frame(
254        &mut self,
255        _cx: &mut Cx,
256        _timeout: Option<std::time::Duration>,
257    ) -> Result<Option<ServerFrame>> {
258        let mut pending = self
259            .pending
260            .lock()
261            .map_err(|_| Error::HostError("registry transport mutex poisoned".to_owned()))?;
262        Ok(pending.take())
263    }
264
265    fn close(&mut self, _cx: &mut Cx) -> Result<()> {
266        let mut pending = self
267            .pending
268            .lock()
269            .map_err(|_| Error::HostError("registry transport mutex poisoned".to_owned()))?;
270        *pending = None;
271        Ok(())
272    }
273
274    fn as_any(&self) -> &dyn std::any::Any {
275        self
276    }
277}
278
279#[cfg(feature = "wasm")]
280pub struct WasmConnectionTransport {
281    region: String,
282    pending: Option<ServerFrame>,
283}
284
285#[cfg(feature = "wasm")]
286impl WasmConnectionTransport {
287    pub fn connect(address: &ServerAddress) -> Result<Self> {
288        let ServerAddress::Wasm { region } = address else {
289            return Err(Error::Eval(
290                "wasm connection transport requires a wasm address".to_owned(),
291            ));
292        };
293        let _ = lookup_wasm_region(region)?;
294        Ok(Self {
295            region: region.clone(),
296            pending: None,
297        })
298    }
299}
300
301#[cfg(feature = "wasm")]
302impl ConnectionTransport for WasmConnectionTransport {
303    fn send_frame(&mut self, _cx: &mut Cx, frame: ServerFrame) -> Result<()> {
304        let region = lookup_wasm_region(&self.region)?;
305        let request = encode_transport_frame(&frame)?;
306        enforce_wasm_transport_limit(&request, "wasm frame exceeds transport limit")?;
307        let reply = region.runtime.call(
308            region.module,
309            &sim_kernel::Symbol::qualified("server", "answer"),
310            WasmFrame::new(request),
311        )?;
312        enforce_wasm_frame_limit(&reply, "wasm reply exceeds transport limit")?;
313        self.pending = Some(decode_transport_frame(reply.bytes())?);
314        Ok(())
315    }
316
317    fn recv_frame(
318        &mut self,
319        _cx: &mut Cx,
320        _timeout: Option<std::time::Duration>,
321    ) -> Result<Option<ServerFrame>> {
322        Ok(self.pending.take())
323    }
324
325    fn close(&mut self, _cx: &mut Cx) -> Result<()> {
326        self.pending = None;
327        Ok(())
328    }
329
330    fn as_any(&self) -> &dyn std::any::Any {
331        self
332    }
333}
334
335#[cfg(feature = "wasm")]
336pub(super) fn enforce_wasm_transport_limit(bytes: &[u8], message: &str) -> Result<()> {
337    if bytes.len() > WasmFrameLimits::default().max_frame_bytes {
338        return Err(Error::HostError(message.to_owned()));
339    }
340    Ok(())
341}
342
343#[cfg(feature = "wasm")]
344pub(super) fn enforce_wasm_frame_limit(frame: &WasmFrame, message: &str) -> Result<()> {
345    let frame_ref = frame.as_ref()?;
346    if usize::try_from(frame_ref.len).unwrap_or(usize::MAX)
347        > WasmFrameLimits::default().max_frame_bytes
348    {
349        return Err(Error::HostError(message.to_owned()));
350    }
351    Ok(())
352}
353
354#[cfg(not(feature = "wasm"))]
355pub struct WasmConnectionTransport;
356
357#[cfg(not(feature = "wasm"))]
358impl WasmConnectionTransport {
359    pub fn connect(_address: &ServerAddress) -> Result<Self> {
360        Err(Error::Eval("server wasm feature disabled".to_owned()))
361    }
362}
363
364#[cfg(not(feature = "wasm"))]
365impl ConnectionTransport for WasmConnectionTransport {
366    fn send_frame(&mut self, _cx: &mut Cx, _frame: ServerFrame) -> Result<()> {
367        Err(Error::Eval("server wasm feature disabled".to_owned()))
368    }
369
370    fn recv_frame(
371        &mut self,
372        _cx: &mut Cx,
373        _timeout: Option<std::time::Duration>,
374    ) -> Result<Option<ServerFrame>> {
375        Err(Error::Eval("server wasm feature disabled".to_owned()))
376    }
377
378    fn close(&mut self, _cx: &mut Cx) -> Result<()> {
379        Ok(())
380    }
381
382    fn as_any(&self) -> &dyn std::any::Any {
383        self
384    }
385}
sim_lib_server/transport/backends.rs

sim_lib_server/transport/
backends.rs
