zakura_network/zakura/transport/
service.rs

1//! Zakura protocol service trait surface.
2
3use std::{collections::HashMap, fmt, future::Future, net::IpAddr, pin::Pin};
4
5use thiserror::Error;
6use tokio_util::sync::CancellationToken;
7
8use super::{CloseCause, FramedRecv, FramedSend};
9use crate::{
10    zakura::{ServicePeerDirection, ZakuraConnId, ZakuraPeerId},
11    BoxError,
12};
13
14use super::Frame;
15
16/// Boxed future returned by object-safe stream handlers.
17pub type BoxRunFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
18
19/// Transport mode for a service-declared stream.
20#[derive(Copy, Clone, Debug, Eq, PartialEq)]
21pub enum StreamMode {
22    /// A long-lived ordered stream between connected peers.
23    Ordered,
24    /// A short-lived request/response stream opened per request.
25    RequestResponse,
26}
27
28/// A service-declared Zakura stream.
29#[derive(Copy, Clone, Debug, Eq, PartialEq)]
30pub struct Stream {
31    /// Unique stream kind carried in `StreamPrelude.stream_kind`.
32    pub kind: u16,
33    /// Version of this stream kind.
34    pub version: u16,
35    /// Maximum application frame bytes for this stream.
36    pub frame_cap: u32,
37    /// Capability bit both peers must negotiate before this stream is wired.
38    pub capability: u64,
39    /// Stream lifetime and opening semantics.
40    pub mode: StreamMode,
41}
42
43/// Transport state for one ordered service stream.
44#[derive(Debug)]
45pub(crate) struct ServiceStream {
46    pub(crate) session_id: u64,
47    pub(crate) recv: FramedRecv,
48    pub(crate) send: FramedSend,
49    pub(crate) cancel_token: CancellationToken,
50}
51
52impl ServiceStream {
53    pub(crate) fn new(
54        session_id: u64,
55        recv: FramedRecv,
56        send: FramedSend,
57        cancel_token: CancellationToken,
58    ) -> Self {
59        Self {
60            session_id,
61            recv,
62            send,
63            cancel_token,
64        }
65    }
66}
67
68/// Per-peer transport state handed to a service when a peer connects.
69#[derive(Debug)]
70pub struct Peer {
71    /// Authenticated Zakura peer identity.
72    pub id: ZakuraPeerId,
73    /// Supervisor registration generation that owns this service session.
74    pub conn_id: ZakuraConnId,
75    /// Remote IP address when the transport knows it.
76    pub remote_ip: Option<IpAddr>,
77    /// Capabilities accepted by both peers.
78    pub negotiated: u64,
79    /// Direction of the underlying authenticated connection.
80    pub direction: ServicePeerDirection,
81    streams: HashMap<u16, ServiceStream>,
82    cancel_token: CancellationToken,
83    service_cancel_token: CancellationToken,
84    close_cause: CloseCause,
85}
86
87impl Peer {
88    /// Build a peer from already-opened transport streams.
89    pub fn new(
90        id: ZakuraPeerId,
91        remote_ip: Option<IpAddr>,
92        negotiated: u64,
93        streams: HashMap<u16, (FramedRecv, FramedSend)>,
94        cancel_token: CancellationToken,
95    ) -> Self {
96        Self::new_with_conn_id_and_direction(
97            0,
98            id,
99            remote_ip,
100            negotiated,
101            ServicePeerDirection::Inbound,
102            streams,
103            cancel_token,
104        )
105    }
106
107    /// Build a peer from already-opened transport streams and a known direction.
108    pub fn new_with_direction(
109        id: ZakuraPeerId,
110        remote_ip: Option<IpAddr>,
111        negotiated: u64,
112        direction: ServicePeerDirection,
113        streams: HashMap<u16, (FramedRecv, FramedSend)>,
114        cancel_token: CancellationToken,
115    ) -> Self {
116        Self::new_with_conn_id_and_direction(
117            0,
118            id,
119            remote_ip,
120            negotiated,
121            direction,
122            streams,
123            cancel_token,
124        )
125    }
126
127    /// Build a peer from transport streams, connection id, and direction.
128    pub(crate) fn new_with_conn_id_and_direction(
129        conn_id: ZakuraConnId,
130        id: ZakuraPeerId,
131        remote_ip: Option<IpAddr>,
132        negotiated: u64,
133        direction: ServicePeerDirection,
134        streams: HashMap<u16, (FramedRecv, FramedSend)>,
135        cancel_token: CancellationToken,
136    ) -> Self {
137        Self::new_with_conn_id_and_direction_and_close_cause(
138            conn_id,
139            id,
140            remote_ip,
141            negotiated,
142            direction,
143            streams,
144            cancel_token,
145            CloseCause::new(),
146        )
147    }
148
149    #[allow(clippy::too_many_arguments)]
150    pub(crate) fn new_with_conn_id_and_direction_and_close_cause(
151        conn_id: ZakuraConnId,
152        id: ZakuraPeerId,
153        remote_ip: Option<IpAddr>,
154        negotiated: u64,
155        direction: ServicePeerDirection,
156        streams: HashMap<u16, (FramedRecv, FramedSend)>,
157        cancel_token: CancellationToken,
158        close_cause: CloseCause,
159    ) -> Self {
160        let streams = streams
161            .into_iter()
162            .map(|(kind, (recv, send))| {
163                (
164                    kind,
165                    ServiceStream::new(0, recv, send, cancel_token.child_token()),
166                )
167            })
168            .collect::<HashMap<_, _>>();
169        Self::new_with_service_streams(
170            conn_id,
171            id,
172            remote_ip,
173            negotiated,
174            direction,
175            streams,
176            cancel_token,
177            close_cause,
178        )
179    }
180
181    #[allow(clippy::too_many_arguments)]
182    pub(crate) fn new_with_service_streams(
183        conn_id: ZakuraConnId,
184        id: ZakuraPeerId,
185        remote_ip: Option<IpAddr>,
186        negotiated: u64,
187        direction: ServicePeerDirection,
188        streams: HashMap<u16, ServiceStream>,
189        cancel_token: CancellationToken,
190        close_cause: CloseCause,
191    ) -> Self {
192        let service_cancel_token = streams
193            .values()
194            .next()
195            .map(|stream| stream.cancel_token.clone())
196            .unwrap_or_else(|| cancel_token.child_token());
197        Self::new_with_service_cancel_token(
198            conn_id,
199            id,
200            remote_ip,
201            negotiated,
202            direction,
203            streams,
204            cancel_token,
205            service_cancel_token,
206            close_cause,
207        )
208    }
209
210    #[allow(clippy::too_many_arguments)]
211    pub(crate) fn new_with_service_cancel_token(
212        conn_id: ZakuraConnId,
213        id: ZakuraPeerId,
214        remote_ip: Option<IpAddr>,
215        negotiated: u64,
216        direction: ServicePeerDirection,
217        streams: HashMap<u16, ServiceStream>,
218        cancel_token: CancellationToken,
219        service_cancel_token: CancellationToken,
220        close_cause: CloseCause,
221    ) -> Self {
222        Self {
223            id,
224            conn_id,
225            remote_ip,
226            negotiated,
227            direction,
228            streams,
229            cancel_token,
230            service_cancel_token,
231            close_cause,
232        }
233    }
234
235    /// Take ownership of a stream pair for `kind`.
236    pub fn take_stream(&mut self, kind: u16) -> Option<(FramedRecv, FramedSend)> {
237        self.streams
238            .remove(&kind)
239            .map(|stream| (stream.recv, stream.send))
240    }
241
242    /// Take ownership of a stream pair and its owning ordered-stream generation.
243    pub fn take_stream_with_session_id(
244        &mut self,
245        kind: u16,
246    ) -> Option<(u64, FramedRecv, FramedSend)> {
247        self.streams
248            .remove(&kind)
249            .map(|stream| (stream.session_id, stream.recv, stream.send))
250    }
251
252    /// Return the cancellation token for this peer's service tasks.
253    ///
254    /// The token is the transport supervisor's per-peer disconnect token, so it
255    /// fires when this peer disconnects or the local node shuts down.
256    pub fn cancel_token(&self) -> CancellationToken {
257        self.cancel_token.clone()
258    }
259
260    /// Return the cancellation token for this service's local session work.
261    ///
262    /// Cancelling this token parks only the local service session. The parent
263    /// peer token still fires on connection shutdown.
264    pub fn service_cancel_token(&self) -> CancellationToken {
265        self.service_cancel_token.clone()
266    }
267
268    /// Return the shared first-cause connection close recorder.
269    pub(crate) fn close_cause(&self) -> CloseCause {
270        self.close_cause.clone()
271    }
272
273    /// Split this peer into fields so the registry can fan streams out by owner.
274    pub(crate) fn into_parts(
275        self,
276    ) -> (
277        ZakuraPeerId,
278        ZakuraConnId,
279        Option<IpAddr>,
280        u64,
281        ServicePeerDirection,
282        HashMap<u16, ServiceStream>,
283        CancellationToken,
284        CloseCause,
285    ) {
286        (
287            self.id,
288            self.conn_id,
289            self.remote_ip,
290            self.negotiated,
291            self.direction,
292            self.streams,
293            self.cancel_token,
294            self.close_cause,
295        )
296    }
297}
298
299/// A Zakura protocol service.
300pub trait Service: fmt::Debug + Send + Sync + 'static {
301    /// Stable service name for logs and diagnostics.
302    fn name(&self) -> &'static str;
303
304    /// Streams this service owns.
305    fn streams(&self) -> &[Stream];
306
307    /// Return whether this service currently wants a new session for `peer`.
308    ///
309    /// This is a cheap, advisory demand check used by the transport before
310    /// opening an ordered stream. Implementations intentionally keep this to
311    /// local room/interest state; remote first-party summary preference is
312    /// applied upstream before dialing or escalation selection.
313    ///
314    /// The service remains authoritative at [`Service::add_peer`], where it
315    /// can still reject or locally park a session if its state changed
316    /// concurrently.
317    fn wants_peer(
318        &self,
319        _peer: &ZakuraPeerId,
320        _negotiated: u64,
321        _direction: ServicePeerDirection,
322    ) -> bool {
323        true
324    }
325
326    /// Add a connected peer and spawn any per-stream work owned by this service.
327    fn add_peer(&self, peer: Peer);
328
329    /// Remove a disconnected peer.
330    fn remove_peer(&self, peer: &ZakuraPeerId, conn_id: ZakuraConnId);
331
332    /// Deliver one request-response frame to this service.
333    fn deliver_frame(
334        &self,
335        _peer_id: ZakuraPeerId,
336        _stream_kind: u16,
337        _frame: Frame,
338    ) -> Result<(), SinkReject> {
339        Err(SinkReject::protocol(
340            "service does not accept inbound frames",
341        ))
342    }
343
344    /// Return this service's request/response handler, if it has one.
345    fn as_request_response(&self) -> Option<&dyn RequestResponseService> {
346        None
347    }
348}
349
350/// A Zakura service that accepts one-shot request/response streams.
351pub trait RequestResponseService: Service {
352    /// Deliver one request-response request frame to this service.
353    ///
354    /// `max_frame_bytes` bounds each encoded outbound frame; `max_message_bytes`
355    /// is the peer's negotiated full-message cap. Responders must size outbound
356    /// frame payloads against the smaller of the two so the peer never receives
357    /// a frame larger than its accepted message cap.
358    fn request_frame<'a>(
359        &'a self,
360        peer_id: ZakuraPeerId,
361        stream_kind: u16,
362        request_id: u64,
363        max_frame_bytes: u32,
364        max_message_bytes: u32,
365        frame: Frame,
366    ) -> BoxRunFuture<'a, Result<Vec<Frame>, SinkReject>>;
367}
368
369/// A per-stream reader owned by a service.
370///
371/// This trait deliberately uses an explicit boxed future instead of native
372/// `async fn` or the `async-trait` crate: stream handlers are intended to be
373/// object-dispatched, and the explicit signature keeps that object safety without
374/// adding another dependency.
375pub trait Sink: Send + 'static {
376    /// Run the reader until the stream closes or is rejected.
377    fn run(self: Box<Self>, recv: FramedRecv) -> BoxRunFuture<'static, Result<(), SinkReject>>;
378}
379
380/// A per-stream writer owned by a service.
381///
382/// P2 services can either run this task shape directly or keep a concrete
383/// typed send handle built from the [`FramedSend`] handed to [`Service::add_peer`].
384///
385/// See [`Sink`] for why this uses an explicit boxed future.
386pub trait Source: Send + 'static {
387    /// Run the writer until the stream closes.
388    fn run(self: Box<Self>, send: FramedSend) -> BoxRunFuture<'static, ()>;
389}
390
391/// Reason a service sink rejected a decoded frame stream.
392#[derive(Debug, Error)]
393pub enum SinkReject {
394    /// The peer sent protocol-invalid data, so the connection should close.
395    #[error("inbound sink rejected protocol-invalid frame: {0}")]
396    Protocol(#[source] BoxError),
397
398    /// Local sink state prevented delivery; the peer is not at fault.
399    #[error("inbound sink could not accept frame locally: {0}")]
400    Local(#[source] BoxError),
401}
402
403impl SinkReject {
404    /// Build a fatal peer-protocol rejection.
405    pub fn protocol(error: impl Into<BoxError>) -> Self {
406        Self::Protocol(error.into())
407    }
408
409    /// Build a non-fatal local-delivery rejection.
410    pub fn local(error: impl Into<BoxError>) -> Self {
411        Self::Local(error.into())
412    }
413}
414
415#[cfg(test)]
416mod tests {
417    use super::*;
418
419    #[test]
420    fn sink_reject_constructors_preserve_protocol_and_local_contract() {
421        let protocol = SinkReject::protocol("bad frame");
422        let local = SinkReject::local("closed queue");
423
424        assert!(matches!(protocol, SinkReject::Protocol(_)));
425        assert!(matches!(local, SinkReject::Local(_)));
426        assert!(protocol.to_string().contains("protocol-invalid"));
427        assert!(local.to_string().contains("locally"));
428    }
429}
zakura_network/zakura/transport/service.rs

zakura_network/zakura/transport/
service.rs
