Skip to main content

kunobi_daemon/wire/
asynchronous.rs

1use super::*;
2use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
3
4/// Async equivalent of [`Session`]. Wrap establishment in one timeout. Cancelling
5/// a frame read or write poisons the session; it cannot resume on a byte prefix.
6pub struct AsyncSession<S> {
7    io: S,
8    buffer: Vec<u8>,
9    cache: buffa::SizeCache,
10    agreement: Agreement,
11    failed: bool,
12}
13
14impl<S: AsyncRead + AsyncWrite + Unpin> AsyncSession<S> {
15    fn unestablished(io: S) -> Self {
16        Self {
17            io,
18            buffer: Vec::new(),
19            cache: buffa::SizeCache::new(),
20            agreement: Agreement {
21                version: VERSION,
22                capabilities: 0,
23                max_frame: HELLO_LIMIT,
24            },
25            failed: false,
26        }
27    }
28
29    /// Establish a client session on an advertised binary endpoint.
30    pub async fn connect(io: S, offer: &Hello) -> io::Result<Self> {
31        offer.validate()?;
32        let mut session = Self::unestablished(io);
33        session.io.write_all(&MAGIC).await?;
34        session.send_message(offer).await?;
35        let remote: Hello = session.receive_message().await?;
36        check_refusal(&remote)?;
37        let agreement = negotiate(offer, &remote)?;
38        session.io.write_all(&[1]).await?;
39        session.io.flush().await?;
40        session.read_acceptance().await?;
41        session.agreement = agreement;
42        Ok(session)
43    }
44
45    /// Establish a server session after the caller has authenticated the peer.
46    pub async fn accept(io: S, offer: &Hello) -> io::Result<Self> {
47        offer.validate()?;
48        let (mut session, remote) = Self::offered(io).await?;
49        let agreement = match negotiate(offer, &remote) {
50            Ok(agreement) => agreement,
51            Err(error) => {
52                if let Some(refusal) = rejection(offer, &remote) {
53                    // The failed negotiation is the error to report.
54                    let _ = session.send_message(&refusal).await;
55                }
56                return Err(error);
57            }
58        };
59        session.send_message(offer).await?;
60        session.read_acceptance().await?;
61        session.io.write_all(&[1]).await?;
62        session.io.flush().await?;
63        session.agreement = agreement;
64        Ok(session)
65    }
66
67    /// Async equivalent of [`Session::refuse`].
68    pub async fn refuse(io: S, offer: &Hello, reason: Refused) -> io::Result<()> {
69        let (mut session, remote) = Self::offered(io).await?;
70        match reason.frame(offer, &remote) {
71            Some(refusal) => session.send_message(&refusal).await,
72            None => Ok(()),
73        }
74    }
75
76    /// Read the magic and the client's offer.
77    async fn offered(io: S) -> io::Result<(Self, Hello)> {
78        let mut session = Self::unestablished(io);
79        let mut magic = [0; 8];
80        session.io.read_exact(&mut magic).await?;
81        if magic != MAGIC {
82            return Err(invalid("wrong binary protocol magic"));
83        }
84        let remote = session.receive_message().await?;
85        Ok((session, remote))
86    }
87
88    async fn read_acceptance(&mut self) -> io::Result<()> {
89        if self.io.read_u8().await? != 1 {
90            return Err(invalid("peer did not accept negotiation"));
91        }
92        Ok(())
93    }
94
95    /// Negotiated capabilities and message bound.
96    pub fn agreement(&self) -> Agreement {
97        self.agreement
98    }
99
100    fn usable(&self) -> io::Result<()> {
101        if self.failed {
102            Err(invalid("session failed; reconnect without replay"))
103        } else {
104            Ok(())
105        }
106    }
107
108    fn check(&self, message: &Control) -> io::Result<()> {
109        if required_capability(message)? & self.agreement.capabilities == 0 {
110            return Err(invalid("operation capability was not negotiated"));
111        }
112        Ok(())
113    }
114
115    /// Send one operation. Failure or cancellation must never trigger replay.
116    pub async fn send(&mut self, message: &Control) -> io::Result<()> {
117        self.check(message)?;
118        self.send_message(message).await
119    }
120
121    /// Receive one operation; cancellation prevents further use of this session.
122    pub async fn receive(&mut self) -> io::Result<Control> {
123        let message = self.receive_message().await?;
124        if let Err(error) = self.check(&message) {
125            self.failed = true;
126            return Err(error);
127        }
128        Ok(message)
129    }
130
131    /// Wait for a subscription peer to disconnect. No further requests are
132    /// legal on that connection. Unlike a framed receive, cancelling this
133    /// single-byte read before completion leaves the session usable.
134    pub(crate) async fn subscription_closed(&mut self) -> io::Result<()> {
135        self.usable()?;
136        let mut byte = [0];
137        let result = self.io.read(&mut byte).await;
138        self.failed = true;
139        match result? {
140            0 => Ok(()),
141            _ => Err(invalid("request after WATCH")),
142        }
143    }
144
145    async fn send_message<M: Message>(&mut self, message: &M) -> io::Result<()> {
146        self.usable()?;
147        encode_frame(
148            message,
149            self.agreement.max_frame,
150            &mut self.buffer,
151            &mut self.cache,
152        )?;
153        self.failed = true;
154        self.io.write_all(&self.buffer).await?;
155        self.io.flush().await?;
156        self.failed = false;
157        Ok(())
158    }
159
160    async fn receive_message<M: Message + Default>(&mut self) -> io::Result<M> {
161        self.usable()?;
162        self.failed = true;
163        let length = self.io.read_u32_le().await?;
164        if length == 0 || length > self.agreement.max_frame {
165            return Err(invalid("received frame exceeds negotiated bounds"));
166        }
167        self.buffer.resize(length as usize, 0);
168        self.io.read_exact(&mut self.buffer).await?;
169        let message = decode_message(self.buffer.as_slice(), self.agreement.max_frame)?;
170        self.failed = false;
171        Ok(message)
172    }
173}