1
//! Declare DataStream, a type that wraps DataReader and DataWriter so as to be useful
2
//! for byte-oriented communication.
3

            
4
use crate::{Error, Result};
5
use static_assertions::assert_impl_all;
6
use tor_cell::relaycell::msg::EndReason;
7
use tor_cell::relaycell::{RelayCellFormat, RelayCmd};
8

            
9
use futures::io::{AsyncRead, AsyncWrite};
10
use futures::stream::StreamExt;
11
use futures::task::{Context, Poll};
12
use futures::{Future, Stream};
13
use pin_project::pin_project;
14
use postage::watch;
15

            
16
#[cfg(feature = "tokio")]
17
use tokio_crate::io::ReadBuf;
18
#[cfg(feature = "tokio")]
19
use tokio_crate::io::{AsyncRead as TokioAsyncRead, AsyncWrite as TokioAsyncWrite};
20
#[cfg(feature = "tokio")]
21
use tokio_util::compat::{FuturesAsyncReadCompatExt, FuturesAsyncWriteCompatExt};
22
use tor_cell::restricted_msg;
23

            
24
use std::fmt::Debug;
25
use std::io::Result as IoResult;
26
use std::num::NonZero;
27
use std::pin::Pin;
28
#[cfg(any(feature = "stream-ctrl", feature = "experimental-api"))]
29
use std::sync::Arc;
30
#[cfg(feature = "stream-ctrl")]
31
use std::sync::{Mutex, Weak};
32

            
33
use educe::Educe;
34

            
35
use crate::client::ClientTunnel;
36
use crate::memquota::StreamAccount;
37
use crate::stream::StreamReceiver;
38
use crate::stream::StreamTarget;
39
use crate::stream::Tunnel;
40
use crate::stream::cmdcheck::{AnyCmdChecker, CmdChecker, StreamStatus};
41
use crate::stream::flow_ctrl::state::StreamRateLimit;
42
use crate::stream::flow_ctrl::xon_xoff::reader::{BufferIsEmpty, XonXoffReader, XonXoffReaderCtrl};
43
use tor_async_utils::rate_limited_writer::{
44
    DynamicRateLimitedWriter, RateLimitedWriter, RateLimitedWriterConfig,
45
};
46
use tor_basic_utils::onionperf_types::{OnionperfEvent, OnionperfStreamStatus};
47
use tor_basic_utils::skip_fmt;
48
use tor_cell::relaycell::msg::Data;
49
use tor_error::internal;
50
use tor_rtcompat::{CoarseTimeProvider, DynTimeProvider, SleepProvider};
51

            
52
/// A stream of [`RateLimitedWriterConfig`] used to update a [`DynamicRateLimitedWriter`].
53
///
54
/// Unfortunately we need to store the result of a [`StreamExt::map`] and [`StreamExt::fuse`] in
55
/// [`DataWriter`], which leaves us with this ugly type.
56
/// We use a type alias to make `DataWriter` a little nicer.
57
type RateConfigStream = futures::stream::Map<
58
    futures::stream::Fuse<watch::Receiver<StreamRateLimit>>,
59
    fn(StreamRateLimit) -> RateLimitedWriterConfig,
60
>;
61

            
62
/// An anonymized stream over the Tor network.
63
///
64
/// For most purposes, you can think of this type as an anonymized
65
/// TCP stream: it can read and write data, and get closed when it's done.
66
///
67
/// [`DataStream`] implements [`futures::io::AsyncRead`] and
68
/// [`futures::io::AsyncWrite`], so you can use it anywhere that those
69
/// traits are expected.
70
///
71
/// # Examples
72
///
73
/// Connecting to an HTTP server and sending a request, using
74
/// [`AsyncWriteExt::write_all`](futures::io::AsyncWriteExt::write_all):
75
///
76
/// ```ignore
77
/// let mut stream = tor_client.connect(("icanhazip.com", 80), None).await?;
78
///
79
/// use futures::io::AsyncWriteExt;
80
///
81
/// stream
82
///     .write_all(b"GET / HTTP/1.1\r\nHost: icanhazip.com\r\nConnection: close\r\n\r\n")
83
///     .await?;
84
///
85
/// // Flushing the stream is important; see below!
86
/// stream.flush().await?;
87
/// ```
88
///
89
/// Reading the result, using [`AsyncReadExt::read_to_end`](futures::io::AsyncReadExt::read_to_end):
90
///
91
/// ```ignore
92
/// use futures::io::AsyncReadExt;
93
///
94
/// let mut buf = Vec::new();
95
/// stream.read_to_end(&mut buf).await?;
96
///
97
/// println!("{}", String::from_utf8_lossy(&buf));
98
/// ```
99
///
100
/// # Usage with Tokio
101
///
102
/// If the `tokio` crate feature is enabled, this type also implements
103
/// [`tokio::io::AsyncRead`](tokio_crate::io::AsyncRead) and
104
/// [`tokio::io::AsyncWrite`](tokio_crate::io::AsyncWrite) for easier integration
105
/// with code that expects those traits.
106
///
107
/// # Remember to call `flush`!
108
///
109
/// DataStream buffers data internally, in order to write as few cells
110
/// as possible onto the network.  In order to make sure that your
111
/// data has actually been sent, you need to make sure that
112
/// [`AsyncWrite::poll_flush`] runs to completion: probably via
113
/// [`AsyncWriteExt::flush`](futures::io::AsyncWriteExt::flush).
114
///
115
/// # Splitting the type
116
///
117
/// This type is internally composed of a [`DataReader`] and a [`DataWriter`]; the
118
/// `DataStream::split` method can be used to split it into those two parts, for more
119
/// convenient usage with e.g. stream combinators.
120
///
121
/// # How long does a stream live?
122
///
123
/// A `DataStream` will live until all references to it are dropped,
124
/// or until it is closed explicitly.
125
///
126
/// If you split the stream into a `DataReader` and a `DataWriter`, it
127
/// will survive until _both_ are dropped, or until it is closed
128
/// explicitly.
129
///
130
/// A stream can also close because of a network error,
131
/// or because the other side of the stream decided to close it.
132
///
133
// # Semver note
134
//
135
// Note that this type is re-exported as a part of the public API of
136
// the `arti-client` crate.  Any changes to its API here in
137
// `tor-proto` need to be reflected above.
138
#[derive(Debug)]
139
pub struct DataStream {
140
    /// Underlying writer for this stream
141
    w: DataWriter,
142
    /// Underlying reader for this stream
143
    r: DataReader,
144
    /// A control object that can be used to monitor and control this stream
145
    /// without needing to own it.
146
    ///
147
    /// Set to `None` if this is not a client stream.
148
    #[cfg(feature = "stream-ctrl")]
149
    ctrl: Option<Arc<ClientDataStreamCtrl>>,
150
}
151
assert_impl_all! { DataStream: Send, Sync }
152

            
153
/// An object used to control and monitor a data stream.
154
///
155
/// # Notes
156
///
157
/// This is a separate type from [`DataStream`] because it's useful to have
158
/// multiple references to this object, whereas a [`DataReader`] and [`DataWriter`]
159
/// need to have a single owner for the `AsyncRead` and `AsyncWrite` APIs to
160
/// work correctly.
161
#[cfg(feature = "stream-ctrl")]
162
#[cfg_attr(
163
    feature = "rpc",
164
    derive(derive_deftly::Deftly),
165
    derive_deftly(tor_rpcbase::templates::Object)
166
)]
167
#[derive(Debug)]
168
pub struct ClientDataStreamCtrl {
169
    /// The circuit to which this stream is attached.
170
    ///
171
    /// Note that the stream's reader and writer halves each contain a `StreamTarget`,
172
    /// which in turn has a strong reference to the `ClientCirc`.  So as long as any
173
    /// one of those is alive, this reference will be present.
174
    ///
175
    /// We make this a Weak reference so that once the stream itself is closed,
176
    /// we can't leak circuits.
177
    tunnel: Weak<ClientTunnel>,
178

            
179
    /// Shared user-visible information about the state of this stream.
180
    ///
181
    /// TODO RPC: This will probably want to be a `postage::Watch` or something
182
    /// similar, if and when it stops moving around.
183
    #[cfg(feature = "stream-ctrl")]
184
    status: Arc<Mutex<DataStreamStatus>>,
185

            
186
    /// The memory quota account that should be used for this stream's data
187
    ///
188
    /// Exists to keep the account alive
189
    _memquota: StreamAccount,
190
}
191

            
192
/// The inner writer for [`DataWriter`].
193
///
194
/// This type is responsible for taking bytes and packaging them into cells.
195
/// Rate limiting is implemented in [`DataWriter`] to avoid making this type more complex.
196
#[derive(Debug)]
197
struct DataWriterInner {
198
    /// Internal state for this writer
199
    ///
200
    /// This is stored in an Option so that we can mutate it in the
201
    /// AsyncWrite functions.  It might be possible to do better here,
202
    /// and we should refactor if so.
203
    state: Option<DataWriterState>,
204

            
205
    /// The memory quota account that should be used for this stream's data
206
    ///
207
    /// Exists to keep the account alive
208
    // If we liked, we could make this conditional; see DataReaderInner.memquota
209
    _memquota: StreamAccount,
210

            
211
    /// A control object that can be used to monitor and control this stream
212
    /// without needing to own it.
213
    ///
214
    /// Set to `None` if this is not a client stream.
215
    #[cfg(feature = "stream-ctrl")]
216
    ctrl: Option<Arc<ClientDataStreamCtrl>>,
217
}
218

            
219
/// The write half of a [`DataStream`], implementing [`futures::io::AsyncWrite`].
220
///
221
/// See the [`DataStream`] docs for more information. In particular, note
222
/// that this writer requires `poll_flush` to complete in order to guarantee that
223
/// all data has been written.
224
///
225
/// # Usage with Tokio
226
///
227
/// If the `tokio` crate feature is enabled, this type also implements
228
/// [`tokio::io::AsyncWrite`](tokio_crate::io::AsyncWrite) for easier integration
229
/// with code that expects that trait.
230
///
231
/// # Drop and close
232
///
233
/// Note that dropping a `DataWriter` has no special effect on its own:
234
/// if the `DataWriter` is dropped, the underlying stream will still remain open
235
/// until the `DataReader` is also dropped.
236
///
237
/// If you want the stream to close earlier, use [`close`](futures::io::AsyncWriteExt::close)
238
/// (or [`shutdown`](tokio_crate::io::AsyncWriteExt::shutdown) with `tokio`).
239
///
240
/// Remember that Tor does not support half-open streams:
241
/// If you `close` or `shutdown` a stream,
242
/// the other side will not see the stream as half-open,
243
/// and so will (probably) not finish sending you any in-progress data.
244
/// Do not use `close`/`shutdown` to communicate anything besides
245
/// "I am done using this stream."
246
///
247
// # Semver note
248
//
249
// Note that this type is re-exported as a part of the public API of
250
// the `arti-client` crate.  Any changes to its API here in
251
// `tor-proto` need to be reflected above.
252
#[derive(Debug)]
253
pub struct DataWriter {
254
    /// A wrapper around [`DataWriterInner`] that adds rate limiting.
255
    writer: DynamicRateLimitedWriter<DataWriterInner, RateConfigStream, DynTimeProvider>,
256
}
257

            
258
impl DataWriter {
259
    /// Create a new rate-limited [`DataWriter`] from a [`DataWriterInner`].
260
124
    fn new(
261
124
        inner: DataWriterInner,
262
124
        rate_limit_updates: watch::Receiver<StreamRateLimit>,
263
124
        time_provider: DynTimeProvider,
264
124
    ) -> Self {
265
        /// Converts a `rate` into a `RateLimitedWriterConfig`.
266
184
        fn rate_to_config(rate: StreamRateLimit) -> RateLimitedWriterConfig {
267
184
            let rate = rate.bytes_per_sec();
268
184
            RateLimitedWriterConfig {
269
184
                rate,        // bytes per second
270
184
                burst: rate, // bytes
271
184
                // This number is chosen arbitrarily, but the idea is that we want to balance
272
184
                // between throughput and latency. Assume the user tries to write a large buffer
273
184
                // (~600 bytes). If we set this too small (for example 1), we'll be waking up
274
184
                // frequently and writing a small number of bytes each time to the
275
184
                // `DataWriterInner`, even if this isn't enough bytes to send a cell. If we set this
276
184
                // too large (for example 510), we'll be waking up infrequently to write a larger
277
184
                // number of bytes each time. So even if the `DataWriterInner` has almost a full
278
184
                // cell's worth of data queued (for example 490) and only needs 509-490=19 more
279
184
                // bytes before a cell can be sent, it will block until the rate limiter allows 510
280
184
                // more bytes.
281
184
                //
282
184
                // TODO(arti#2028): Is there an optimal value here?
283
184
                wake_when_bytes_available: NonZero::new(200).expect("200 != 0"), // bytes
284
184
            }
285
184
        }
286

            
287
        // get the current rate from the `watch::Receiver`, which we'll use as the initial rate
288
124
        let initial_rate: StreamRateLimit = *rate_limit_updates.borrow();
289

            
290
        // map the rate update stream to the type required by `DynamicRateLimitedWriter`
291
124
        let rate_limit_updates = rate_limit_updates.fuse().map(rate_to_config as fn(_) -> _);
292

            
293
        // build the rate limiter
294
124
        let writer = RateLimitedWriter::new(inner, &rate_to_config(initial_rate), time_provider);
295
124
        let writer = DynamicRateLimitedWriter::new(writer, rate_limit_updates);
296

            
297
124
        Self { writer }
298
124
    }
299

            
300
    /// Return a [`ClientDataStreamCtrl`] object that can be used to monitor and
301
    /// interact with this stream without holding the stream itself.
302
    ///
303
    /// Returns `None` if this is not a client stream.
304
    #[cfg(feature = "stream-ctrl")]
305
    pub fn client_stream_ctrl(&self) -> Option<&Arc<ClientDataStreamCtrl>> {
306
        self.writer.inner().client_stream_ctrl()
307
    }
308
}
309

            
310
impl AsyncWrite for DataWriter {
311
5540
    fn poll_write(
312
5540
        mut self: Pin<&mut Self>,
313
5540
        cx: &mut Context<'_>,
314
5540
        buf: &[u8],
315
5540
    ) -> Poll<IoResult<usize>> {
316
5540
        AsyncWrite::poll_write(Pin::new(&mut self.writer), cx, buf)
317
5540
    }
318

            
319
32
    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<IoResult<()>> {
320
32
        AsyncWrite::poll_flush(Pin::new(&mut self.writer), cx)
321
32
    }
322

            
323
16
    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<IoResult<()>> {
324
16
        AsyncWrite::poll_close(Pin::new(&mut self.writer), cx)
325
16
    }
326
}
327

            
328
#[cfg(feature = "tokio")]
329
impl TokioAsyncWrite for DataWriter {
330
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<IoResult<usize>> {
331
        TokioAsyncWrite::poll_write(Pin::new(&mut self.compat_write()), cx, buf)
332
    }
333

            
334
    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<IoResult<()>> {
335
        TokioAsyncWrite::poll_flush(Pin::new(&mut self.compat_write()), cx)
336
    }
337

            
338
    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<IoResult<()>> {
339
        TokioAsyncWrite::poll_shutdown(Pin::new(&mut self.compat_write()), cx)
340
    }
341
}
342

            
343
/// The read half of a [`DataStream`], implementing [`futures::io::AsyncRead`].
344
///
345
/// See the [`DataStream`] docs for more information.
346
///
347
/// # Usage with Tokio
348
///
349
/// If the `tokio` crate feature is enabled, this type also implements
350
/// [`tokio::io::AsyncRead`](tokio_crate::io::AsyncRead) for easier integration
351
/// with code that expects that trait.
352
//
353
// # Semver note
354
//
355
// Note that this type is re-exported as a part of the public API of
356
// the `arti-client` crate.  Any changes to its API here in
357
// `tor-proto` need to be reflected above.
358
#[derive(Debug)]
359
pub struct DataReader {
360
    /// The [`DataReaderInner`] with a wrapper to support XON/XOFF flow control.
361
    reader: XonXoffReader<DataReaderInner>,
362
}
363

            
364
impl DataReader {
365
    /// Create a new [`DataReader`].
366
124
    fn new(reader: DataReaderInner, xon_xoff_reader_ctrl: XonXoffReaderCtrl) -> Self {
367
124
        Self {
368
124
            reader: XonXoffReader::new(xon_xoff_reader_ctrl, reader),
369
124
        }
370
124
    }
371

            
372
    /// Return a [`ClientDataStreamCtrl`] object that can be used to monitor and
373
    /// interact with this stream without holding the stream itself.
374
    ///
375
    /// Returns `None` if this is not a client stream.
376
    #[cfg(feature = "stream-ctrl")]
377
    pub fn client_stream_ctrl(&self) -> Option<&Arc<ClientDataStreamCtrl>> {
378
        self.reader.inner().client_stream_ctrl()
379
    }
380
}
381

            
382
impl AsyncRead for DataReader {
383
244
    fn poll_read(
384
244
        mut self: Pin<&mut Self>,
385
244
        cx: &mut Context<'_>,
386
244
        buf: &mut [u8],
387
244
    ) -> Poll<IoResult<usize>> {
388
244
        AsyncRead::poll_read(Pin::new(&mut self.reader), cx, buf)
389
244
    }
390

            
391
    fn poll_read_vectored(
392
        mut self: Pin<&mut Self>,
393
        cx: &mut Context<'_>,
394
        bufs: &mut [std::io::IoSliceMut<'_>],
395
    ) -> Poll<IoResult<usize>> {
396
        AsyncRead::poll_read_vectored(Pin::new(&mut self.reader), cx, bufs)
397
    }
398
}
399

            
400
#[cfg(feature = "tokio")]
401
impl TokioAsyncRead for DataReader {
402
    fn poll_read(
403
        self: Pin<&mut Self>,
404
        cx: &mut Context<'_>,
405
        buf: &mut ReadBuf<'_>,
406
    ) -> Poll<IoResult<()>> {
407
        TokioAsyncRead::poll_read(Pin::new(&mut self.compat()), cx, buf)
408
    }
409
}
410

            
411
/// The inner reader for [`DataReader`].
412
///
413
/// This type is responsible for taking stream messages and extracting the stream data from them.
414
/// Flow control logic is implemented in [`DataReader`] to avoid making this type more complex.
415
#[derive(Debug)]
416
pub(crate) struct DataReaderInner {
417
    /// Internal state for this reader.
418
    ///
419
    /// This is stored in an Option so that we can mutate it in
420
    /// poll_read().  It might be possible to do better here, and we
421
    /// should refactor if so.
422
    state: Option<DataReaderState>,
423

            
424
    /// The memory quota account that should be used for this stream's data
425
    ///
426
    /// Exists to keep the account alive
427
    // If we liked, we could make this conditional on not(cfg(feature = "stream-ctrl"))
428
    // since, ClientDataStreamCtrl contains a StreamAccount clone too.  But that seems fragile.
429
    _memquota: StreamAccount,
430

            
431
    /// A control object that can be used to monitor and control this stream
432
    /// without needing to own it.
433
    ///
434
    /// Set to `None` if this is not a client stream.
435
    #[cfg(feature = "stream-ctrl")]
436
    ctrl: Option<Arc<ClientDataStreamCtrl>>,
437
}
438

            
439
impl BufferIsEmpty for DataReaderInner {
440
    /// The result will become stale,
441
    /// so is most accurate immediately after a [`poll_read`](AsyncRead::poll_read).
442
    fn is_empty(mut self: Pin<&mut Self>) -> bool {
443
        match self
444
            .state
445
            .as_mut()
446
            .expect("forgot to put `DataReaderState` back")
447
        {
448
            DataReaderState::Open(imp) => {
449
                // check if the partial cell in `pending` is empty,
450
                // and if the message stream is empty
451
                imp.pending[imp.offset..].is_empty() && imp.s.is_empty()
452
            }
453
            // closed, so any data should have been discarded
454
            DataReaderState::Closed => true,
455
        }
456
    }
457
}
458

            
459
/// Shared status flags for tracking the status of as `DataStream`.
460
///
461
/// We expect to refactor this a bit, so it's not exposed at all.
462
//
463
// TODO RPC: Possibly instead of manipulating the fields of DataStreamStatus
464
// from various points in this module, we should instead construct
465
// DataStreamStatus as needed from information available elsewhere.  In any
466
// case, we should really  eliminate as much duplicate state here as we can.
467
// (See discussions at !1198 for some challenges with this.)
468
#[cfg(feature = "stream-ctrl")]
469
#[derive(Clone, Debug, Default)]
470
struct DataStreamStatus {
471
    /// True if we've received a CONNECTED message.
472
    //
473
    // TODO: This is redundant with `connected` in DataReaderImpl.
474
    received_connected: bool,
475
    /// True if we have decided to send an END message.
476
    //
477
    // TODO RPC: There is not an easy way to set this from this module!  Really,
478
    // the decision to send an "end" is made when the StreamTarget object is
479
    // dropped, but we don't currently have any way to see when that happens.
480
    // Perhaps we need a different shared StreamStatus object that the
481
    // StreamTarget holds?
482
    sent_end: bool,
483
    /// True if we have received an END message telling us to close the stream.
484
    received_end: bool,
485
    /// True if we have received an error.
486
    ///
487
    /// (This is not a subset or superset of received_end; some errors are END
488
    /// messages but some aren't; some END messages are errors but some aren't.)
489
    received_err: bool,
490
}
491

            
492
#[cfg(feature = "stream-ctrl")]
493
impl DataStreamStatus {
494
    /// Remember that we've received a connected message.
495
124
    fn record_connected(&mut self) {
496
124
        self.received_connected = true;
497
124
    }
498

            
499
    /// Remember that we've received an error of some kind.
500
24
    fn record_error(&mut self, e: &Error) {
501
        // TODO: Probably we should remember the actual error in a box or
502
        // something.  But that means making a redundant copy of the error
503
        // even if nobody will want it.  Do we care?
504
24
        match e {
505
24
            Error::EndReceived(EndReason::DONE) => self.received_end = true,
506
            Error::EndReceived(_) => {
507
                self.received_end = true;
508
                self.received_err = true;
509
            }
510
            _ => self.received_err = true,
511
        }
512
24
    }
513
}
514

            
515
restricted_msg! {
516
    /// An allowable incoming message on a client data stream.
517
    enum ClientDataStreamMsg:RelayMsg {
518
        // SENDME is handled by the reactor.
519
        Data, End, Connected,
520
    }
521
}
522

            
523
// TODO RPC: Should we also implement this trait for everything that holds a
524
// ClientDataStreamCtrl?
525
#[cfg(feature = "stream-ctrl")]
526
impl super::ctrl::ClientStreamCtrl for ClientDataStreamCtrl {
527
    fn tunnel(&self) -> Option<Arc<ClientTunnel>> {
528
        self.tunnel.upgrade()
529
    }
530
}
531

            
532
#[cfg(feature = "stream-ctrl")]
533
impl ClientDataStreamCtrl {
534
    /// Return true if the underlying stream is connected. (That is, if it has
535
    /// received a `CONNECTED` message, and has not been closed.)
536
    pub fn is_connected(&self) -> bool {
537
        let s = self.status.lock().expect("poisoned lock");
538
        s.received_connected && !(s.sent_end || s.received_end || s.received_err)
539
    }
540

            
541
    // TODO RPC: Add more functions once we have the desired API more nailed
542
    // down.
543
}
544

            
545
impl DataStream {
546
    /// Wrap raw stream receiver and target parts as a DataStream.
547
    ///
548
    /// For non-optimistic stream, function `wait_for_connection`
549
    /// must be called after to make sure CONNECTED is received.
550
96
    pub(crate) fn new<P: SleepProvider + CoarseTimeProvider>(
551
96
        time_provider: P,
552
96
        receiver: StreamReceiver,
553
96
        xon_xoff_reader_ctrl: XonXoffReaderCtrl,
554
96
        target: StreamTarget,
555
96
        memquota: StreamAccount,
556
96
    ) -> Self {
557
96
        Self::new_inner(
558
96
            time_provider,
559
96
            receiver,
560
96
            xon_xoff_reader_ctrl,
561
96
            target,
562
            false,
563
96
            memquota,
564
        )
565
96
    }
566

            
567
    /// Wrap raw stream receiver and target parts as a connected DataStream.
568
    ///
569
    /// Unlike [`DataStream::new`], this creates a `DataStream` that does not expect to receive a
570
    /// CONNECTED cell.
571
    ///
572
    /// This is used by hidden services, exit relays, and directory servers to accept streams.
573
    #[cfg(any(feature = "hs-service", feature = "relay"))]
574
28
    pub(crate) fn new_connected<P: SleepProvider + CoarseTimeProvider>(
575
28
        time_provider: P,
576
28
        receiver: StreamReceiver,
577
28
        xon_xoff_reader_ctrl: XonXoffReaderCtrl,
578
28
        target: StreamTarget,
579
28
        memquota: StreamAccount,
580
28
    ) -> Self {
581
28
        Self::new_inner(
582
28
            time_provider,
583
28
            receiver,
584
28
            xon_xoff_reader_ctrl,
585
28
            target,
586
            true,
587
28
            memquota,
588
        )
589
28
    }
590

            
591
    /// The shared implementation of the `new*()` functions.
592
124
    fn new_inner<P: SleepProvider + CoarseTimeProvider>(
593
124
        time_provider: P,
594
124
        receiver: StreamReceiver,
595
124
        xon_xoff_reader_ctrl: XonXoffReaderCtrl,
596
124
        target: StreamTarget,
597
124
        connected: bool,
598
124
        memquota: StreamAccount,
599
124
    ) -> Self {
600
124
        let relay_cell_format = target.relay_cell_format();
601
124
        let out_buf_len = Data::max_body_len(relay_cell_format);
602
124
        let rate_limit_stream = target.rate_limit_stream().clone();
603

            
604
124
        tracing::trace!(
605
            onionperf = true,
606
            event = ?OnionperfEvent::Stream(OnionperfStreamStatus::New),
607
            stream_id = ?target.stream_id,
608
            circ_id = ?match target.clone().tunnel {
609
                Tunnel::Client(client) => Some(client.circ.unique_id()),
610
                #[cfg(feature = "relay")]
611
                Tunnel::Relay(_) => None, // TODO
612
            },
613
        );
614

            
615
        #[cfg(feature = "stream-ctrl")]
616
124
        let status = {
617
124
            let mut data_stream_status = DataStreamStatus::default();
618
124
            if connected {
619
28
                data_stream_status.record_connected();
620
96
            }
621
124
            Arc::new(Mutex::new(data_stream_status))
622
        };
623

            
624
        #[cfg(feature = "stream-ctrl")]
625
124
        let ctrl = {
626
124
            let tunnel = match target.tunnel() {
627
120
                crate::stream::Tunnel::Client(t) => Some(Arc::downgrade(t)),
628
                #[cfg(feature = "relay")]
629
4
                crate::stream::Tunnel::Relay(_) => None,
630
            };
631

            
632
124
            tunnel.map(|tunnel| {
633
120
                Arc::new(ClientDataStreamCtrl {
634
120
                    tunnel,
635
120
                    status: status.clone(),
636
120
                    _memquota: memquota.clone(),
637
120
                })
638
120
            })
639
        };
640
124
        let r = DataReaderInner {
641
124
            state: Some(DataReaderState::Open(DataReaderImpl {
642
124
                s: receiver,
643
124
                pending: Vec::new(),
644
124
                offset: 0,
645
124
                connected,
646
124
                #[cfg(feature = "stream-ctrl")]
647
124
                status: status.clone(),
648
124
            })),
649
124
            _memquota: memquota.clone(),
650
124
            #[cfg(feature = "stream-ctrl")]
651
124
            ctrl: ctrl.clone(),
652
124
        };
653
124
        let w = DataWriterInner {
654
124
            state: Some(DataWriterState::Ready(DataWriterImpl {
655
124
                s: target,
656
124
                buf: vec![0; out_buf_len].into_boxed_slice(),
657
124
                n_pending: 0,
658
124
                #[cfg(feature = "stream-ctrl")]
659
124
                status,
660
124
                relay_cell_format,
661
124
            })),
662
124
            _memquota: memquota,
663
124
            #[cfg(feature = "stream-ctrl")]
664
124
            ctrl: ctrl.clone(),
665
124
        };
666

            
667
124
        let time_provider = DynTimeProvider::new(time_provider);
668

            
669
124
        DataStream {
670
124
            w: DataWriter::new(w, rate_limit_stream, time_provider),
671
124
            r: DataReader::new(r, xon_xoff_reader_ctrl),
672
124
            #[cfg(feature = "stream-ctrl")]
673
124
            ctrl,
674
124
        }
675
124
    }
676

            
677
    /// Divide this DataStream into its constituent parts.
678
28
    pub fn split(self) -> (DataReader, DataWriter) {
679
28
        (self.r, self.w)
680
28
    }
681

            
682
    /// Wait until a CONNECTED cell is received, or some other cell
683
    /// is received to indicate an error.
684
    ///
685
    /// Does nothing if this stream is already connected.
686
126
    pub async fn wait_for_connection(&mut self) -> Result<()> {
687
        // We must put state back before returning
688
84
        let state = self
689
84
            .r
690
84
            .reader
691
84
            .inner_mut()
692
84
            .state
693
84
            .take()
694
84
            .expect("Missing state in DataReaderInner");
695

            
696
84
        if let DataReaderState::Open(mut imp) = state {
697
84
            let result = if imp.connected {
698
                Ok(())
699
            } else {
700
                // This succeeds if the cell is CONNECTED, and fails otherwise.
701
188
                std::future::poll_fn(|cx| Pin::new(&mut imp).read_cell(cx)).await
702
            };
703
84
            self.r.reader.inner_mut().state = Some(match result {
704
                Err(_) => DataReaderState::Closed,
705
84
                Ok(_) => DataReaderState::Open(imp),
706
            });
707
84
            result
708
        } else {
709
            Err(Error::from(internal!(
710
                "Expected ready state, got {:?}",
711
                state
712
            )))
713
        }
714
84
    }
715

            
716
    /// Return a [`ClientDataStreamCtrl`] object that can be used to monitor and
717
    /// interact with this stream without holding the stream itself.
718
    #[cfg(feature = "stream-ctrl")]
719
    pub fn client_stream_ctrl(&self) -> Option<&Arc<ClientDataStreamCtrl>> {
720
        self.ctrl.as_ref()
721
    }
722
}
723

            
724
impl AsyncRead for DataStream {
725
244
    fn poll_read(
726
244
        mut self: Pin<&mut Self>,
727
244
        cx: &mut Context<'_>,
728
244
        buf: &mut [u8],
729
244
    ) -> Poll<IoResult<usize>> {
730
244
        AsyncRead::poll_read(Pin::new(&mut self.r), cx, buf)
731
244
    }
732
}
733

            
734
#[cfg(feature = "tokio")]
735
impl TokioAsyncRead for DataStream {
736
    fn poll_read(
737
        self: Pin<&mut Self>,
738
        cx: &mut Context<'_>,
739
        buf: &mut ReadBuf<'_>,
740
    ) -> Poll<IoResult<()>> {
741
        TokioAsyncRead::poll_read(Pin::new(&mut self.compat()), cx, buf)
742
    }
743
}
744

            
745
impl AsyncWrite for DataStream {
746
5540
    fn poll_write(
747
5540
        mut self: Pin<&mut Self>,
748
5540
        cx: &mut Context<'_>,
749
5540
        buf: &[u8],
750
5540
    ) -> Poll<IoResult<usize>> {
751
5540
        AsyncWrite::poll_write(Pin::new(&mut self.w), cx, buf)
752
5540
    }
753
32
    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<IoResult<()>> {
754
32
        AsyncWrite::poll_flush(Pin::new(&mut self.w), cx)
755
32
    }
756
    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<IoResult<()>> {
757
        AsyncWrite::poll_close(Pin::new(&mut self.w), cx)
758
    }
759
}
760

            
761
#[cfg(feature = "tokio")]
762
impl TokioAsyncWrite for DataStream {
763
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<IoResult<usize>> {
764
        TokioAsyncWrite::poll_write(Pin::new(&mut self.compat()), cx, buf)
765
    }
766

            
767
    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<IoResult<()>> {
768
        TokioAsyncWrite::poll_flush(Pin::new(&mut self.compat()), cx)
769
    }
770

            
771
    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<IoResult<()>> {
772
        TokioAsyncWrite::poll_shutdown(Pin::new(&mut self.compat()), cx)
773
    }
774
}
775

            
776
/// Helper type: Like BoxFuture, but also requires that the future be Sync.
777
type BoxSyncFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + Sync + 'a>>;
778

            
779
/// An enumeration for the state of a DataWriter.
780
///
781
/// We have to use an enum here because, for as long as we're waiting
782
/// for a flush operation to complete, the future returned by
783
/// `flush_cell()` owns the DataWriterImpl.
784
#[derive(Educe)]
785
#[educe(Debug)]
786
enum DataWriterState {
787
    /// The writer has closed or gotten an error: nothing more to do.
788
    Closed,
789
    /// The writer is not currently flushing; more data can get queued
790
    /// immediately.
791
    Ready(DataWriterImpl),
792
    /// The writer is flushing a cell.
793
    Flushing(
794
        #[educe(Debug(method = "skip_fmt"))] //
795
        BoxSyncFuture<'static, (DataWriterImpl, Result<()>)>,
796
    ),
797
}
798

            
799
/// Internal: the write part of a DataStream
800
#[derive(Educe)]
801
#[educe(Debug)]
802
struct DataWriterImpl {
803
    /// The underlying StreamTarget object.
804
    s: StreamTarget,
805

            
806
    /// Buffered data to send over the connection.
807
    ///
808
    /// This buffer is currently allocated using a number of bytes
809
    /// equal to the maximum that we can package at a time.
810
    //
811
    // TODO: this buffer is probably smaller than we want, but it's good
812
    // enough for now.  If we _do_ make it bigger, we'll have to change
813
    // our use of Data::split_from to handle the case where we can't fit
814
    // all the data.
815
    #[educe(Debug(method = "skip_fmt"))]
816
    buf: Box<[u8]>,
817

            
818
    /// Number of unflushed bytes in buf.
819
    n_pending: usize,
820

            
821
    /// Relay cell format in use
822
    relay_cell_format: RelayCellFormat,
823

            
824
    /// Shared user-visible information about the state of this stream.
825
    #[cfg(feature = "stream-ctrl")]
826
    status: Arc<Mutex<DataStreamStatus>>,
827
}
828

            
829
impl DataWriterInner {
830
    /// See [`DataWriter::client_stream_ctrl`].
831
    #[cfg(feature = "stream-ctrl")]
832
    fn client_stream_ctrl(&self) -> Option<&Arc<ClientDataStreamCtrl>> {
833
        self.ctrl.as_ref()
834
    }
835

            
836
    /// Helper for poll_flush() and poll_close(): Performs a flush, then
837
    /// closes the stream if should_close is true.
838
48
    fn poll_flush_impl(
839
48
        mut self: Pin<&mut Self>,
840
48
        cx: &mut Context<'_>,
841
48
        should_close: bool,
842
48
    ) -> Poll<IoResult<()>> {
843
48
        let state = self.state.take().expect("Missing state in DataWriter");
844

            
845
        // TODO: this whole function is a bit copy-pasted.
846
48
        let mut future: BoxSyncFuture<_> = match state {
847
48
            DataWriterState::Ready(imp) => {
848
48
                if imp.n_pending == 0 {
849
                    // Nothing to flush!
850
24
                    if should_close {
851
                        // We need to actually continue with this function to do the closing.
852
                        // Thus, make a future that does nothing and is ready immediately.
853
16
                        Box::pin(futures::future::ready((imp, Ok(()))))
854
                    } else {
855
                        // There's nothing more to do; we can return.
856
8
                        self.state = Some(DataWriterState::Ready(imp));
857
8
                        return Poll::Ready(Ok(()));
858
                    }
859
                } else {
860
                    // We need to flush the buffer's contents; Make a future for that.
861
24
                    Box::pin(imp.flush_buf())
862
                }
863
            }
864
            DataWriterState::Flushing(fut) => fut,
865
            DataWriterState::Closed => {
866
                self.state = Some(DataWriterState::Closed);
867
                return Poll::Ready(Err(Error::NotConnected.into()));
868
            }
869
        };
870

            
871
40
        match future.as_mut().poll(cx) {
872
            Poll::Ready((imp, Err(e))) => {
873
                match e {
874
                    Error::NotConnected => (),
875
                    _ => tracing::trace!(
876
                        onionperf = true,
877
                        event = ?OnionperfEvent::Stream(OnionperfStreamStatus::Failed),
878
                        stream_id = ?imp.s.stream_id,
879
                        reason = ?e,
880
                    ),
881
                }
882
                self.state = Some(DataWriterState::Closed);
883
                Poll::Ready(Err(e.into()))
884
            }
885
40
            Poll::Ready((mut imp, Ok(()))) => {
886
40
                if should_close {
887
                    // Tell the StreamTarget to close, so that the reactor
888
                    // realizes that we are done sending. (Dropping `imp.s` does not
889
                    // suffice, since there may be other clones of it.  In particular,
890
                    // the StreamReceiver has one, which it uses to keep the stream
891
                    // open, among other things.)
892
16
                    imp.s.close();
893

            
894
                    #[cfg(feature = "stream-ctrl")]
895
16
                    {
896
16
                        // TODO RPC:  This is not sufficient to track every case
897
16
                        // where we might have sent an End.  See note on the
898
16
                        // `sent_end` field.
899
16
                        imp.status.lock().expect("lock poisoned").sent_end = true;
900
16
                    }
901
16
                    tracing::trace!(
902
                        onionperf = true,
903
                        event = ?OnionperfEvent::Stream(OnionperfStreamStatus::Closed),
904
                        stream_id = ?imp.s.stream_id,
905
                    );
906
16
                    self.state = Some(DataWriterState::Closed);
907
24
                } else {
908
24
                    self.state = Some(DataWriterState::Ready(imp));
909
24
                }
910
40
                Poll::Ready(Ok(()))
911
            }
912
            Poll::Pending => {
913
                self.state = Some(DataWriterState::Flushing(future));
914
                Poll::Pending
915
            }
916
        }
917
48
    }
918
}
919

            
920
impl AsyncWrite for DataWriterInner {
921
5540
    fn poll_write(
922
5540
        mut self: Pin<&mut Self>,
923
5540
        cx: &mut Context<'_>,
924
5540
        buf: &[u8],
925
5540
    ) -> Poll<IoResult<usize>> {
926
5540
        if buf.is_empty() {
927
            return Poll::Ready(Ok(0));
928
5540
        }
929

            
930
5540
        let state = self.state.take().expect("Missing state in DataWriter");
931

            
932
5540
        let mut future = match state {
933
5500
            DataWriterState::Ready(mut imp) => {
934
5500
                let n_queued = imp.queue_bytes(buf);
935
5500
                if n_queued != 0 {
936
1220
                    self.state = Some(DataWriterState::Ready(imp));
937
1220
                    return Poll::Ready(Ok(n_queued));
938
4280
                }
939
                // we couldn't queue anything, so the current cell must be full.
940
4280
                Box::pin(imp.flush_buf())
941
            }
942
40
            DataWriterState::Flushing(fut) => fut,
943
            DataWriterState::Closed => {
944
                self.state = Some(DataWriterState::Closed);
945
                return Poll::Ready(Err(Error::NotConnected.into()));
946
            }
947
        };
948

            
949
4320
        match future.as_mut().poll(cx) {
950
            Poll::Ready((_imp, Err(e))) => {
951
                #[cfg(feature = "stream-ctrl")]
952
                {
953
                    _imp.status.lock().expect("lock poisoned").record_error(&e);
954
                }
955
                self.state = Some(DataWriterState::Closed);
956
                Poll::Ready(Err(e.into()))
957
            }
958
4280
            Poll::Ready((mut imp, Ok(()))) => {
959
                // Great!  We're done flushing.  Queue as much as we can of this
960
                // cell.
961
4280
                let n_queued = imp.queue_bytes(buf);
962
4280
                self.state = Some(DataWriterState::Ready(imp));
963
4280
                Poll::Ready(Ok(n_queued))
964
            }
965
            Poll::Pending => {
966
40
                self.state = Some(DataWriterState::Flushing(future));
967
40
                Poll::Pending
968
            }
969
        }
970
5540
    }
971

            
972
32
    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<IoResult<()>> {
973
32
        self.poll_flush_impl(cx, false)
974
32
    }
975

            
976
16
    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<IoResult<()>> {
977
16
        self.poll_flush_impl(cx, true)
978
16
    }
979
}
980

            
981
#[cfg(feature = "tokio")]
982
impl TokioAsyncWrite for DataWriterInner {
983
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<IoResult<usize>> {
984
        TokioAsyncWrite::poll_write(Pin::new(&mut self.compat_write()), cx, buf)
985
    }
986

            
987
    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<IoResult<()>> {
988
        TokioAsyncWrite::poll_flush(Pin::new(&mut self.compat_write()), cx)
989
    }
990

            
991
    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<IoResult<()>> {
992
        TokioAsyncWrite::poll_shutdown(Pin::new(&mut self.compat_write()), cx)
993
    }
994
}
995

            
996
impl DataWriterImpl {
997
    /// Try to flush the current buffer contents as a data cell.
998
6456
    async fn flush_buf(mut self) -> (Self, Result<()>) {
999
4304
        let result = if let Some((cell, remainder)) =
4304
            Data::try_split_from(self.relay_cell_format, &self.buf[..self.n_pending])
        {
            // TODO: Eventually we may want a larger buffer; if we do,
            // this invariant will become false.
4304
            assert!(remainder.is_empty());
4304
            self.n_pending = 0;
4304
            self.s.send(cell.into()).await
        } else {
            Ok(())
        };
4304
        (self, result)
4304
    }
    /// Add as many bytes as possible from `b` to our internal buffer;
    /// return the number we were able to add.
9780
    fn queue_bytes(&mut self, b: &[u8]) -> usize {
9780
        let empty_space = &mut self.buf[self.n_pending..];
9780
        if empty_space.is_empty() {
            // that is, len == 0
4280
            return 0;
5500
        }
5500
        let n_to_copy = std::cmp::min(b.len(), empty_space.len());
5500
        empty_space[..n_to_copy].copy_from_slice(&b[..n_to_copy]);
5500
        self.n_pending += n_to_copy;
5500
        n_to_copy
9780
    }
}
impl DataReaderInner {
    /// Return a [`ClientDataStreamCtrl`] object that can be used to monitor and
    /// interact with this stream without holding the stream itself.
    #[cfg(feature = "stream-ctrl")]
    pub(crate) fn client_stream_ctrl(&self) -> Option<&Arc<ClientDataStreamCtrl>> {
        self.ctrl.as_ref()
    }
}
/// An enumeration for the state of a [`DataReaderInner`].
// TODO: We don't need to implement the state in this way anymore now that we've removed the saved
// future. There are a few ways we could simplify this. See:
// https://gitlab.torproject.org/tpo/core/arti/-/merge_requests/3076#note_3218210
#[derive(Educe)]
#[educe(Debug)]
// We allow this since it's expected that streams will spend most of their time in the `Open` state,
// and will be cleaned up shortly after closing.
#[allow(clippy::large_enum_variant)]
enum DataReaderState {
    /// In this state we have received an end cell or an error.
    Closed,
    /// In this state the reader is open.
    Open(DataReaderImpl),
}
/// Wrapper for the read part of a [`DataStream`].
#[derive(Educe)]
#[educe(Debug)]
#[pin_project]
struct DataReaderImpl {
    /// The underlying StreamReceiver object.
    #[educe(Debug(method = "skip_fmt"))]
    #[pin]
    s: StreamReceiver,
    /// If present, data that we received on this stream but have not
    /// been able to send to the caller yet.
    // TODO: This data structure is probably not what we want, but
    // it's good enough for now.
    #[educe(Debug(method = "skip_fmt"))]
    pending: Vec<u8>,
    /// Index into pending to show what we've already read.
    offset: usize,
    /// If true, we have received a CONNECTED cell on this stream.
    connected: bool,
    /// Shared user-visible information about the state of this stream.
    #[cfg(feature = "stream-ctrl")]
    status: Arc<Mutex<DataStreamStatus>>,
}
impl AsyncRead for DataReaderInner {
244
    fn poll_read(
244
        mut self: Pin<&mut Self>,
244
        cx: &mut Context<'_>,
244
        buf: &mut [u8],
244
    ) -> Poll<IoResult<usize>> {
        // We're pulling the state object out of the reader.  We MUST
        // put it back before this function returns.
244
        let mut state = self.state.take().expect("Missing state in DataReaderInner");
        loop {
348
            let mut imp = match state {
348
                DataReaderState::Open(mut imp) => {
                    // There may be data to read already.
348
                    let n_copied = imp.extract_bytes(buf);
348
                    if n_copied != 0 || buf.is_empty() {
                        // We read data into the buffer, or the buffer was 0-len to begin with.
                        // Tell the caller.
92
                        self.state = Some(DataReaderState::Open(imp));
92
                        return Poll::Ready(Ok(n_copied));
256
                    }
                    // No data available!  We have to try reading.
256
                    imp
                }
                DataReaderState::Closed => {
                    // TODO: Why are we returning an error rather than continuing to return EOF?
                    self.state = Some(DataReaderState::Closed);
                    return Poll::Ready(Err(Error::NotConnected.into()));
                }
            };
            // See if a cell is ready.
256
            match Pin::new(&mut imp).read_cell(cx) {
24
                Poll::Ready(Err(e)) => {
                    // There aren't any survivable errors in the current
                    // design.
24
                    self.state = Some(DataReaderState::Closed);
                    #[cfg(feature = "stream-ctrl")]
24
                    {
24
                        imp.status.lock().expect("lock poisoned").record_error(&e);
24
                    }
24
                    let result = if matches!(e, Error::EndReceived(EndReason::DONE)) {
24
                        Ok(0)
                    } else {
                        Err(e.into())
                    };
24
                    return Poll::Ready(result);
                }
104
                Poll::Ready(Ok(())) => {
104
                    // It read a cell!  Continue the loop.
104
                    state = DataReaderState::Open(imp);
104
                }
                Poll::Pending => {
                    // No cells ready, so tell the
                    // caller to get back to us later.
128
                    self.state = Some(DataReaderState::Open(imp));
128
                    return Poll::Pending;
                }
            }
        }
244
    }
}
#[cfg(feature = "tokio")]
impl TokioAsyncRead for DataReaderInner {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<IoResult<()>> {
        TokioAsyncRead::poll_read(Pin::new(&mut self.compat()), cx, buf)
    }
}
impl DataReaderImpl {
    /// Pull as many bytes as we can off of self.pending, and return that
    /// number of bytes.
348
    fn extract_bytes(&mut self, buf: &mut [u8]) -> usize {
348
        let remainder = &self.pending[self.offset..];
348
        let n_to_copy = std::cmp::min(buf.len(), remainder.len());
348
        buf[..n_to_copy].copy_from_slice(&remainder[..n_to_copy]);
348
        self.offset += n_to_copy;
348
        n_to_copy
348
    }
    /// Return true iff there are no buffered bytes here to yield
92
    fn buf_is_empty(&self) -> bool {
92
        self.pending.len() == self.offset
92
    }
    /// Load self.pending with the contents of a new data cell.
444
    fn read_cell(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
        use ClientDataStreamMsg::*;
444
        let msg = match self.as_mut().project().s.poll_next(cx) {
232
            Poll::Pending => return Poll::Pending,
212
            Poll::Ready(Some(Ok(unparsed))) => match unparsed.decode::<ClientDataStreamMsg>() {
212
                Ok(cell) => cell.into_msg(),
                Err(e) => {
                    self.s.protocol_error();
                    return Poll::Ready(Err(Error::from_bytes_err(e, "message on a data stream")));
                }
            },
            Poll::Ready(Some(Err(e))) => return Poll::Ready(Err(e)),
            // TODO: This doesn't seem right to me, but seems to be the behaviour of the code before
            // the refactoring, so I've kept the same behaviour. I think if the cell stream is
            // terminated, we should be returning `None` here and not considering it as an error.
            // The `StreamReceiver` will have already returned an error if the cell stream was
            // terminated without an END message.
            Poll::Ready(None) => return Poll::Ready(Err(Error::NotConnected)),
        };
212
        let result = match msg {
96
            Connected(_) if !self.connected => {
96
                self.connected = true;
                #[cfg(feature = "stream-ctrl")]
96
                {
96
                    self.status
96
                        .lock()
96
                        .expect("poisoned lock")
96
                        .record_connected();
96
                }
96
                Ok(())
            }
            Connected(_) => {
                self.s.protocol_error();
                Err(Error::StreamProto(
                    "Received a second connect cell on a data stream".to_string(),
                ))
            }
92
            Data(d) if self.connected => {
92
                self.add_data(d.into());
92
                Ok(())
            }
            Data(_) => {
                self.s.protocol_error();
                Err(Error::StreamProto(
                    "Received a data cell an unconnected stream".to_string(),
                ))
            }
24
            End(e) => Err(Error::EndReceived(e.reason())),
        };
212
        Poll::Ready(result)
444
    }
    /// Add the data from `d` to the end of our pending bytes.
92
    fn add_data(&mut self, mut d: Vec<u8>) {
92
        if self.buf_is_empty() {
92
            // No data pending?  Just take d as the new pending.
92
            self.pending = d;
92
            self.offset = 0;
92
        } else {
            // TODO(nickm) This has potential to grow `pending` without bound.
            // Fortunately, we don't currently read cells or call this
            // `add_data` method when pending is nonempty—but if we do in the
            // future, we'll have to be careful here.
            self.pending.append(&mut d);
        }
92
    }
}
/// A `CmdChecker` that enforces invariants for outbound data streams.
#[derive(Debug)]
pub(crate) struct OutboundDataCmdChecker {
    /// True if we are expecting to receive a CONNECTED message on this stream.
    expecting_connected: bool,
}
impl Default for OutboundDataCmdChecker {
362
    fn default() -> Self {
362
        Self {
362
            expecting_connected: true,
362
        }
362
    }
}
impl CmdChecker for OutboundDataCmdChecker {
1200
    fn check_msg(&mut self, msg: &tor_cell::relaycell::UnparsedRelayMsg) -> Result<StreamStatus> {
        use StreamStatus::*;
1200
        match msg.cmd() {
            RelayCmd::CONNECTED => {
106
                if !self.expecting_connected {
4
                    Err(Error::StreamProto(
4
                        "Received CONNECTED twice on a stream.".into(),
4
                    ))
                } else {
102
                    self.expecting_connected = false;
102
                    Ok(Open)
                }
            }
            RelayCmd::DATA => {
1068
                if !self.expecting_connected {
1068
                    Ok(Open)
                } else {
                    Err(Error::StreamProto(
                        "Received DATA before CONNECTED on a stream".into(),
                    ))
                }
            }
24
            RelayCmd::END => Ok(Closed),
2
            _ => Err(Error::StreamProto(format!(
2
                "Unexpected {} on a data stream!",
2
                msg.cmd()
2
            ))),
        }
1200
    }
1008
    fn consume_checked_msg(&mut self, msg: tor_cell::relaycell::UnparsedRelayMsg) -> Result<()> {
1008
        let _ = msg
1008
            .decode::<ClientDataStreamMsg>()
1008
            .map_err(|err| Error::from_bytes_err(err, "cell on half-closed stream"))?;
1008
        Ok(())
1008
    }
}
impl OutboundDataCmdChecker {
    /// Return a new boxed `DataCmdChecker` in a state suitable for a newly
    /// constructed connection.
362
    pub(crate) fn new_any() -> AnyCmdChecker {
362
        Box::<Self>::default()
362
    }
}