tor_proto/stream/
queue.rs

1//! Queues for stream messages.
2//!
3//! While these are technically "channels", we call them "queues" to indicate that they're mostly
4//! just dumb pipes. They do some tracking (memquota and size), but nothing else. The higher-level
5//! object is [`StreamReceiver`](crate::stream::raw::StreamReceiver) which tracks SENDME and END
6//! messages. So the idea is that the "queue" (ex: [`StreamQueueReceiver`]) just holds data and the
7//! "channel" (ex: `StreamReceiver`) adds the Tor logic.
8//!
9//! The main purpose of these types is so that we can count how many bytes of stream data are
10//! stored for the stream. Ideally we'd use a channel type that tracks and reports this as part of
11//! its implementation, but popular channel implementations don't seem to do that.
12
13use std::fmt::Debug;
14use std::pin::Pin;
15use std::sync::{Arc, Mutex};
16use std::task::{Context, Poll};
17
18use futures::{Sink, SinkExt, Stream};
19use tor_async_utils::SinkTrySend;
20use tor_async_utils::peekable_stream::UnobtrusivePeekableStream;
21use tor_async_utils::stream_peek::StreamUnobtrusivePeeker;
22use tor_cell::relaycell::UnparsedRelayMsg;
23use tor_memquota::mq_queue::{self, ChannelSpec, MpscSpec};
24use tor_rtcompat::DynTimeProvider;
25
26use crate::memquota::{SpecificAccount, StreamAccount};
27
28/// Create a new stream queue for incoming messages
29/// (messages arriving on the stream from the Tor network).
30pub(crate) fn stream_queue(
31    size: usize,
32    memquota: &StreamAccount,
33    time_prov: &DynTimeProvider,
34) -> Result<(StreamQueueSender, StreamQueueReceiver), tor_memquota::Error> {
35    // Note that the size here may be very large,
36    // for example when used with XON/XOFF flow control.
37    //
38    // Someday if we remove support for window-based flow control
39    // and only support XON/XOFF flow control,
40    // we may want to make this unbounded instead.
41    // https://gitlab.torproject.org/tpo/core/arti/-/work_items/2412
42    let (sender, receiver) =
43        MpscSpec::new(size).new_mq(time_prov.clone(), memquota.as_raw_account())?;
44
45    let receiver = StreamUnobtrusivePeeker::new(receiver);
46    let counter = Arc::new(Mutex::new(0));
47    Ok((
48        StreamQueueSender {
49            sender,
50            counter: Arc::clone(&counter),
51        },
52        StreamQueueReceiver { receiver, counter },
53    ))
54}
55
56/// For testing purposes, create a stream queue with a no-op memquota account and a fake time
57/// provider.
58#[cfg(test)]
59pub(crate) fn fake_stream_queue(size: usize) -> (StreamQueueSender, StreamQueueReceiver) {
60    // The fake Account doesn't care about the data ages, so this will do.
61    //
62    // This would be wrong to use generally in tests, where we might want to mock time,
63    // since we end up, here with totally *different* mocked time.
64    // But it's OK here, and saves passing a runtime parameter into this function.
65    stream_queue(
66        size,
67        &StreamAccount::new_noop(),
68        &DynTimeProvider::new(tor_rtmock::MockRuntime::default()),
69    )
70    .expect("create fake stream queue")
71}
72
73/// The sending end of a channel of incoming stream messages.
74#[derive(Debug)]
75#[pin_project::pin_project]
76pub(crate) struct StreamQueueSender {
77    /// The inner sender.
78    #[pin]
79    sender: mq_queue::Sender<UnparsedRelayMsg, MpscSpec>,
80    /// Number of bytes within the queue.
81    counter: Arc<Mutex<usize>>,
82}
83
84/// The receiving end of a channel of incoming stream messages.
85#[derive(Debug)]
86#[pin_project::pin_project]
87pub(crate) struct StreamQueueReceiver {
88    /// The inner receiver.
89    ///
90    /// We add the [`StreamUnobtrusivePeeker`] here so that peeked messages are included in
91    /// `counter`.
92    // TODO(arti#534): the possible extra msg held by the `StreamUnobtrusivePeeker` isn't tracked by
93    // memquota
94    #[pin]
95    receiver: StreamUnobtrusivePeeker<mq_queue::Receiver<UnparsedRelayMsg, MpscSpec>>,
96    /// Number of bytes within the queue.
97    counter: Arc<Mutex<usize>>,
98}
99
100impl StreamQueueSender {
101    /// Get the approximate number of data bytes queued for this stream.
102    ///
103    /// As messages can be dequeued at any time, the return value may be larger than the actual
104    /// number of bytes queued for this stream.
105    pub(crate) fn approx_stream_bytes(&self) -> usize {
106        *self.counter.lock().expect("poisoned")
107    }
108}
109
110impl StreamQueueReceiver {
111    /// Get the approximate number of data bytes queued for this stream.
112    ///
113    /// As messages can be enqueued at any time, the return value may be smaller than the actual
114    /// number of bytes queued for this stream.
115    pub(crate) fn approx_stream_bytes(&self) -> usize {
116        *self.counter.lock().expect("poisoned")
117    }
118}
119
120impl Sink<UnparsedRelayMsg> for StreamQueueSender {
121    type Error = <mq_queue::Sender<UnparsedRelayMsg, MpscSpec> as Sink<UnparsedRelayMsg>>::Error;
122
123    fn poll_ready(
124        mut self: Pin<&mut Self>,
125        cx: &mut Context<'_>,
126    ) -> Poll<std::result::Result<(), Self::Error>> {
127        self.sender.poll_ready_unpin(cx)
128    }
129
130    fn start_send(
131        mut self: Pin<&mut Self>,
132        item: UnparsedRelayMsg,
133    ) -> std::result::Result<(), Self::Error> {
134        let mut self_ = self.as_mut().project();
135
136        let stream_data_len = data_len(&item);
137
138        // This lock ensures that us sending the item and the counter increase are done
139        // "atomically", so that the receiver doesn't see the item and try to decrement the
140        // counter before we've incremented the counter, which could cause an underflow.
141        let mut counter = self_.counter.lock().expect("poisoned");
142
143        self_.sender.start_send_unpin(item)?;
144
145        *counter = counter
146            .checked_add(stream_data_len.into())
147            .expect("queue has more than `usize::MAX` bytes?!");
148
149        Ok(())
150    }
151
152    fn poll_flush(
153        mut self: Pin<&mut Self>,
154        cx: &mut Context<'_>,
155    ) -> Poll<std::result::Result<(), Self::Error>> {
156        self.sender.poll_flush_unpin(cx)
157    }
158
159    fn poll_close(
160        mut self: Pin<&mut Self>,
161        cx: &mut Context<'_>,
162    ) -> Poll<std::result::Result<(), Self::Error>> {
163        self.sender.poll_close_unpin(cx)
164    }
165}
166
167impl SinkTrySend<UnparsedRelayMsg> for StreamQueueSender {
168    type Error =
169        <mq_queue::Sender<UnparsedRelayMsg, MpscSpec> as SinkTrySend<UnparsedRelayMsg>>::Error;
170
171    fn try_send_or_return(
172        mut self: Pin<&mut Self>,
173        item: UnparsedRelayMsg,
174    ) -> Result<
175        (),
176        (
177            <Self as SinkTrySend<UnparsedRelayMsg>>::Error,
178            UnparsedRelayMsg,
179        ),
180    > {
181        let self_ = self.as_mut().project();
182
183        let stream_data_len = data_len(&item);
184
185        // See comments in `StreamQueueSender::start_send`.
186        let mut counter = self_.counter.lock().expect("poisoned");
187
188        self_.sender.try_send_or_return(item)?;
189
190        *counter = counter
191            .checked_add(stream_data_len.into())
192            .expect("queue has more than `usize::MAX` bytes?!");
193
194        Ok(())
195    }
196}
197
198impl Stream for StreamQueueReceiver {
199    type Item = UnparsedRelayMsg;
200
201    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
202        let self_ = self.as_mut().project();
203
204        // This lock ensures that us receiving the item and the counter decrease are done
205        // "atomically", so that the sender doesn't send a new item and try to increase the
206        // counter before we've decreased the counter, which could cause an overflow.
207        let mut counter = self_.counter.lock().expect("poisoned");
208
209        let item = match self_.receiver.poll_next(cx) {
210            Poll::Ready(Some(x)) => x,
211            Poll::Ready(None) => return Poll::Ready(None),
212            Poll::Pending => return Poll::Pending,
213        };
214
215        let stream_data_len = data_len(&item);
216
217        if stream_data_len != 0 {
218            *counter = counter
219                .checked_sub(stream_data_len.into())
220                .expect("we've removed more bytes than we've added?!");
221        }
222
223        Poll::Ready(Some(item))
224    }
225}
226
227impl UnobtrusivePeekableStream for StreamQueueReceiver {
228    fn unobtrusive_peek_mut<'s>(
229        self: Pin<&'s mut Self>,
230    ) -> Option<&'s mut <Self as futures::Stream>::Item> {
231        self.project().receiver.unobtrusive_peek_mut()
232    }
233}
234
235/// The `length` field of the message, or 0 if not a data message.
236///
237/// If the RELAY_DATA message had an invalid length field, we just ignore the message.
238/// The receiver will find out eventually when it tries to parse the message.
239/// We could return an error here, but for now I think it's best not to behave as if this
240/// queue is performing any validation.
241///
242/// This is its own function so that all parts of the code use the same logic.
243fn data_len(item: &UnparsedRelayMsg) -> u16 {
244    item.data_len().unwrap_or(0)
245}
tor_proto/stream/queue.rs

tor_proto/stream/
queue.rs
