1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
//! Declare DataStream, a type that wraps RawCellStream so as to be useful
//! for byte-oriented communication.

use crate::{Error, Result};
use tor_cell::relaycell::msg::EndReason;
use tor_cell::relaycell::RelayCmd;

use futures::io::{AsyncRead, AsyncWrite};
use futures::task::{Context, Poll};
use futures::Future;

#[cfg(feature = "tokio")]
use tokio_crate::io::ReadBuf;
#[cfg(feature = "tokio")]
use tokio_crate::io::{AsyncRead as TokioAsyncRead, AsyncWrite as TokioAsyncWrite};
#[cfg(feature = "tokio")]
use tokio_util::compat::{FuturesAsyncReadCompatExt, FuturesAsyncWriteCompatExt};
use tor_cell::restricted_msg;

use std::fmt::Debug;
use std::io::Result as IoResult;
use std::pin::Pin;
#[cfg(any(feature = "stream-ctrl", feature = "experimental-api"))]
use std::sync::Arc;
#[cfg(feature = "stream-ctrl")]
use std::sync::{Mutex, Weak};

use educe::Educe;

#[cfg(feature = "experimental-api")]
use crate::circuit::ClientCirc;

use crate::circuit::StreamTarget;
use crate::stream::StreamReader;
use tor_basic_utils::skip_fmt;
use tor_cell::relaycell::msg::Data;
use tor_error::internal;

use super::AnyCmdChecker;

/// An anonymized stream over the Tor network.
///
/// For most purposes, you can think of this type as an anonymized
/// TCP stream: it can read and write data, and get closed when it's done.
///
/// [`DataStream`] implements [`futures::io::AsyncRead`] and
/// [`futures::io::AsyncWrite`], so you can use it anywhere that those
/// traits are expected.
///
/// # Examples
///
/// Connecting to an HTTP server and sending a request, using
/// [`AsyncWriteExt::write_all`](futures::io::AsyncWriteExt::write_all):
///
/// ```ignore
/// let mut stream = tor_client.connect(("icanhazip.com", 80), None).await?;
///
/// use futures::io::AsyncWriteExt;
///
/// stream
///     .write_all(b"GET / HTTP/1.1\r\nHost: icanhazip.com\r\nConnection: close\r\n\r\n")
///     .await?;
///
/// // Flushing the stream is important; see below!
/// stream.flush().await?;
/// ```
///
/// Reading the result, using [`AsyncReadExt::read_to_end`](futures::io::AsyncReadExt::read_to_end):
///
/// ```ignore
/// use futures::io::AsyncReadExt;
///
/// let mut buf = Vec::new();
/// stream.read_to_end(&mut buf).await?;
///
/// println!("{}", String::from_utf8_lossy(&buf));
/// ```
///
/// # Usage with Tokio
///
/// If the `tokio` crate feature is enabled, this type also implements
/// [`tokio::io::AsyncRead`](tokio_crate::io::AsyncRead) and
/// [`tokio::io::AsyncWrite`](tokio_crate::io::AsyncWrite) for easier integration
/// with code that expects those traits.
///
/// # Remember to call `flush`!
///
/// DataStream buffers data internally, in order to write as few cells
/// as possible onto the network.  In order to make sure that your
/// data has actually been sent, you need to make sure that
/// [`AsyncWrite::poll_flush`] runs to completion: probably via
/// [`AsyncWriteExt::flush`](futures::io::AsyncWriteExt::flush).
///
/// # Splitting the type
///
/// This type is internally composed of a [`DataReader`] and a [`DataWriter`]; the
/// `DataStream::split` method can be used to split it into those two parts, for more
/// convenient usage with e.g. stream combinators.
// # Semver note
//
// Note that this type is re-exported as a part of the public API of
// the `arti-client` crate.  Any changes to its API here in
// `tor-proto` need to be reflected above.
#[derive(Debug)]
pub struct DataStream {
    /// Underlying writer for this stream
    w: DataWriter,
    /// Underlying reader for this stream
    r: DataReader,
    /// A handle for the underlying circuit.
    ///
    /// TODO: This is redundant with the reference in `StreamTarget` inside
    /// DataWriterState, but for now we can't actually access that state all the time,
    /// since it might be inside a boxed future.
    ///
    /// TODO: This is also redundant with the reference in `DataStreamCtrl`.
    #[cfg(feature = "experimental-api")]
    circuit: Arc<ClientCirc>,

    /// A control object that can be used to monitor and control this stream
    /// without needing to own it.
    #[cfg(feature = "stream-ctrl")]
    ctrl: std::sync::Arc<DataStreamCtrl>,
}

/// An object used to control and monitor a data stream.
///
/// # Notes
///
/// This is a separate type from [`DataStream`] because it's useful to have
/// multiple references to this object, whereas a [`DataReader`] and [`DataWriter`]
/// need to have a single owner for the `AsyncRead` and `AsyncWrite` APIs to
/// work correctly.
#[cfg(feature = "stream-ctrl")]
#[derive(Debug)]
pub struct DataStreamCtrl {
    /// The circuit to which this stream is attached.
    ///
    /// Note that the stream's reader and writer halves each contain a `StreamTarget`,
    /// which in turn has a strong reference to the `ClientCirc`.  So as long as any
    /// one of those is alive, this reference will be present.
    ///
    /// We make this a Weak reference so that once the stream itself is closed,
    /// we can't leak circuits.
    circuit: Weak<ClientCirc>,

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

/// The write half of a [`DataStream`], implementing [`futures::io::AsyncWrite`].
///
/// See the [`DataStream`] docs for more information. In particular, note
/// that this writer requires `poll_flush` to complete in order to guarantee that
/// all data has been written.
///
/// # Usage with Tokio
///
/// If the `tokio` crate feature is enabled, this type also implements
/// [`tokio::io::AsyncWrite`](tokio_crate::io::AsyncWrite) for easier integration
/// with code that expects that trait.
// # Semver note
//
// Note that this type is re-exported as a part of the public API of
// the `arti-client` crate.  Any changes to its API here in
// `tor-proto` need to be reflected above.
#[derive(Debug)]
pub struct DataWriter {
    /// Internal state for this writer
    ///
    /// This is stored in an Option so that we can mutate it in the
    /// AsyncWrite functions.  It might be possible to do better here,
    /// and we should refactor if so.
    state: Option<DataWriterState>,

    /// A control object that can be used to monitor and control this stream
    /// without needing to own it.
    #[cfg(feature = "stream-ctrl")]
    ctrl: std::sync::Arc<DataStreamCtrl>,
}

/// The read half of a [`DataStream`], implementing [`futures::io::AsyncRead`].
///
/// See the [`DataStream`] docs for more information.
///
/// # Usage with Tokio
///
/// If the `tokio` crate feature is enabled, this type also implements
/// [`tokio::io::AsyncRead`](tokio_crate::io::AsyncRead) for easier integration
/// with code that expects that trait.
//
// # Semver note
//
// Note that this type is re-exported as a part of the public API of
// the `arti-client` crate.  Any changes to its API here in
// `tor-proto` need to be reflected above.
#[derive(Debug)]
pub struct DataReader {
    /// Internal state for this reader.
    ///
    /// This is stored in an Option so that we can mutate it in
    /// poll_read().  It might be possible to do better here, and we
    /// should refactor if so.
    state: Option<DataReaderState>,

    /// A control object that can be used to monitor and control this stream
    /// without needing to own it.
    #[cfg(feature = "stream-ctrl")]
    ctrl: std::sync::Arc<DataStreamCtrl>,
}

/// Shared status flags for tracking the status of as `DataStream`.
///
/// We expect to refactor this a bit, so it's not exposed at all.
//
// TODO RPC: Possibly instead of manipulating the fields of DataStreamStatus
// from various points in this module, we should instead construct
// DataStreamStatus as needed from information available elsewhere.  In any
// case, we should really  eliminate as much duplicate state here as we can.
// (See discussions at !1198 for some challenges with this.)
#[cfg(feature = "stream-ctrl")]
#[derive(Clone, Debug, Default)]
struct DataStreamStatus {
    /// True if we've received a CONNECTED message.
    //
    // TODO: This is redundant with `connected` in DataReaderImpl and
    // `expecting_connected` in DataCmdChecker.
    received_connected: bool,
    /// True if we have decided to send an END message.
    //
    // TODO RPC: There is not an easy way to set this from this module!  Really,
    // the decision to send an "end" is made when the StreamTarget object is
    // dropped, but we don't currently have any way to see when that happens.
    // Perhaps we need a different shared StreamStatus object that the
    // StreamTarget holds?
    sent_end: bool,
    /// True if we have received an END message telling us to close the stream.
    received_end: bool,
    /// True if we have received an error.
    ///
    /// (This is not a subset or superset of received_end; some errors are END
    /// messages but some aren't; some END messages are errors but some aren't.)
    received_err: bool,
}

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

    /// Remember that we've received an error of some kind.
    fn record_error(&mut self, e: &Error) {
        // TODO: Probably we should remember the actual error in a box or
        // something.  But that means making a redundant copy of the error
        // even if nobody will want it.  Do we care?
        match e {
            Error::EndReceived(EndReason::DONE) => self.received_end = true,
            Error::EndReceived(_) => {
                self.received_end = true;
                self.received_err = true;
            }
            _ => self.received_err = true,
        }
    }
}

restricted_msg! {
    /// An allowable incoming message on a data stream.
    enum DataStreamMsg:RelayMsg {
        // SENDME is handled by the reactor.
        Data, End, Connected,
    }
}

// TODO RPC: Should we also implement this trait for everything that holds a
// DataStreamCtrl?
#[cfg(feature = "stream-ctrl")]
impl super::ctrl::ClientStreamCtrl for DataStreamCtrl {
    fn circuit(&self) -> Option<Arc<ClientCirc>> {
        self.circuit.upgrade()
    }
}

#[cfg(feature = "stream-ctrl")]
impl DataStreamCtrl {
    /// Return true if the underlying stream is open. (That is, if it has
    /// received a `CONNECTED` message, and has not been closed.)
    //
    // TODO RPC: Maybe this method belongs in ClientStreamCtrl; maybe others do
    // as well!  We need to talk about moving them around.
    pub fn is_open(&self) -> bool {
        let s = self.status.lock().expect("poisoned lock");
        s.received_connected && !(s.sent_end || s.received_end || s.received_err)
    }

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

impl DataStream {
    /// Wrap raw stream reader and target parts as a DataStream.
    ///
    /// For non-optimistic stream, function `wait_for_connection`
    /// must be called after to make sure CONNECTED is received.
    pub(crate) fn new(reader: StreamReader, target: StreamTarget) -> Self {
        Self::new_inner(reader, target, false)
    }

    /// Wrap raw stream reader and target parts as a connected DataStream.
    ///
    /// Unlike [`DataStream::new`], this creates a `DataStream` that does not expect to receive a
    /// CONNECTED cell.
    ///
    /// This is used by hidden services, exit relays, and directory servers to accept streams.
    #[cfg(feature = "hs-service")]
    pub(crate) fn new_connected(reader: StreamReader, target: StreamTarget) -> Self {
        Self::new_inner(reader, target, true)
    }

    /// The shared implementation of the `new*()` functions.
    fn new_inner(reader: StreamReader, target: StreamTarget, connected: bool) -> Self {
        #[cfg(feature = "stream-ctrl")]
        let status = {
            let mut data_stream_status = DataStreamStatus::default();
            if connected {
                data_stream_status.record_connected();
            }
            Arc::new(Mutex::new(data_stream_status))
        };
        #[cfg(feature = "stream-ctrl")]
        let ctrl = Arc::new(DataStreamCtrl {
            circuit: Arc::downgrade(target.circuit()),
            status: status.clone(),
        });
        #[cfg(feature = "experimental-api")]
        let circuit = target.circuit().clone();
        let r = DataReader {
            state: Some(DataReaderState::Ready(DataReaderImpl {
                s: reader,
                pending: Vec::new(),
                offset: 0,
                connected,
                #[cfg(feature = "stream-ctrl")]
                status: status.clone(),
            })),
            #[cfg(feature = "stream-ctrl")]
            ctrl: ctrl.clone(),
        };
        let w = DataWriter {
            state: Some(DataWriterState::Ready(DataWriterImpl {
                s: target,
                buf: Box::new([0; Data::MAXLEN]),
                n_pending: 0,
                #[cfg(feature = "stream-ctrl")]
                status,
            })),
            #[cfg(feature = "stream-ctrl")]
            ctrl: ctrl.clone(),
        };
        DataStream {
            w,
            r,
            #[cfg(feature = "experimental-api")]
            circuit,
            #[cfg(feature = "stream-ctrl")]
            ctrl,
        }
    }

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

    /// Wait until a CONNECTED cell is received, or some other cell
    /// is received to indicate an error.
    ///
    /// Does nothing if this stream is already connected.
    pub(crate) async fn wait_for_connection(&mut self) -> Result<()> {
        // We must put state back before returning
        let state = self.r.state.take().expect("Missing state in DataReader");

        if let DataReaderState::Ready(imp) = state {
            let (imp, result) = if imp.connected {
                (imp, Ok(()))
            } else {
                // This succeeds if the cell is CONNECTED, and fails otherwise.
                imp.read_cell().await
            };
            self.r.state = Some(match result {
                Err(_) => DataReaderState::Closed,
                Ok(_) => DataReaderState::Ready(imp),
            });
            result
        } else {
            Err(Error::from(internal!(
                "Expected ready state, got {:?}",
                state
            )))
        }
    }

    /// Return a reference to this stream's circuit?
    ///
    /// This is an experimental API; it is not covered by semver guarantee. It
    /// is likely to change or disappear in a future release.
    ///
    /// TODO: Should there be an AttachedToCircuit trait that we use for all
    /// client stream types?  Should this return an Option<&ClientCirc>?
    ///
    /// TODO RPC: Perhaps we should deprecate this in favor of `stream.ctrl().circuit()`.
    #[cfg(feature = "experimental-api")]
    pub fn circuit(&self) -> &ClientCirc {
        &self.circuit
    }

    /// Return a [`DataStreamCtrl`] object that can be used to monitor and
    /// interact with this stream without holding the stream itself.
    #[cfg(feature = "stream-ctrl")]
    pub fn ctrl(&self) -> &Arc<DataStreamCtrl> {
        &self.ctrl
    }
}

impl AsyncRead for DataStream {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<IoResult<usize>> {
        AsyncRead::poll_read(Pin::new(&mut self.r), cx, buf)
    }
}

#[cfg(feature = "tokio")]
impl TokioAsyncRead for DataStream {
    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 AsyncWrite for DataStream {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<IoResult<usize>> {
        AsyncWrite::poll_write(Pin::new(&mut self.w), cx, buf)
    }
    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<IoResult<()>> {
        AsyncWrite::poll_flush(Pin::new(&mut self.w), cx)
    }
    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<IoResult<()>> {
        AsyncWrite::poll_close(Pin::new(&mut self.w), cx)
    }
}

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

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

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

/// An enumeration for the state of a DataWriter.
///
/// We have to use an enum here because, for as long as we're waiting
/// for a flush operation to complete, the future returned by
/// `flush_cell()` owns the DataWriterImpl.
#[derive(Educe)]
#[educe(Debug)]
enum DataWriterState {
    /// The writer has closed or gotten an error: nothing more to do.
    Closed,
    /// The writer is not currently flushing; more data can get queued
    /// immediately.
    Ready(DataWriterImpl),
    /// The writer is flushing a cell.
    Flushing(
        #[educe(Debug(method = "skip_fmt"))]
        Pin<Box<dyn Future<Output = (DataWriterImpl, Result<()>)> + Send>>,
    ),
}

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

    /// Buffered data to send over the connection.
    // TODO: this buffer is probably smaller than we want, but it's good
    // enough for now.  If we _do_ make it bigger, we'll have to change
    // our use of Data::split_from to handle the case where we can't fit
    // all the data.
    #[educe(Debug(method = "skip_fmt"))]
    buf: Box<[u8; Data::MAXLEN]>,

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

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

impl DataWriter {
    /// Return a [`DataStreamCtrl`] object that can be used to monitor and
    /// interact with this stream without holding the stream itself.
    #[cfg(feature = "stream-ctrl")]
    pub fn ctrl(&self) -> &Arc<DataStreamCtrl> {
        &self.ctrl
    }

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

        // TODO: this whole function is a bit copy-pasted.

        let mut future = match state {
            DataWriterState::Ready(imp) => {
                if imp.n_pending == 0 {
                    // Nothing to flush!
                    self.state = Some(DataWriterState::Ready(imp));
                    return Poll::Ready(Ok(()));
                }

                Box::pin(imp.flush_buf())
            }
            DataWriterState::Flushing(fut) => fut,
            DataWriterState::Closed => {
                self.state = Some(DataWriterState::Closed);
                return Poll::Ready(Err(Error::NotConnected.into()));
            }
        };

        match future.as_mut().poll(cx) {
            Poll::Ready((_imp, Err(e))) => {
                self.state = Some(DataWriterState::Closed);
                Poll::Ready(Err(e.into()))
            }
            Poll::Ready((imp, Ok(()))) => {
                if should_close {
                    #[cfg(feature = "stream-ctrl")]
                    {
                        // TODO RPC:  This is not sufficient to track every case
                        // where we might have sent an End.  See note on the
                        // `sent_end` field.
                        imp.status.lock().expect("lock poisoned").sent_end = true;
                    }
                    self.state = Some(DataWriterState::Closed);
                } else {
                    self.state = Some(DataWriterState::Ready(imp));
                }
                Poll::Ready(Ok(()))
            }
            Poll::Pending => {
                self.state = Some(DataWriterState::Flushing(future));
                Poll::Pending
            }
        }
    }
}

impl AsyncWrite for DataWriter {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<IoResult<usize>> {
        if buf.is_empty() {
            return Poll::Ready(Ok(0));
        }

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

        let mut future = match state {
            DataWriterState::Ready(mut imp) => {
                let n_queued = imp.queue_bytes(buf);
                if n_queued != 0 {
                    self.state = Some(DataWriterState::Ready(imp));
                    return Poll::Ready(Ok(n_queued));
                }
                // we couldn't queue anything, so the current cell must be full.
                Box::pin(imp.flush_buf())
            }
            DataWriterState::Flushing(fut) => fut,
            DataWriterState::Closed => {
                self.state = Some(DataWriterState::Closed);
                return Poll::Ready(Err(Error::NotConnected.into()));
            }
        };

        match future.as_mut().poll(cx) {
            Poll::Ready((_imp, Err(e))) => {
                #[cfg(feature = "stream-ctrl")]
                {
                    _imp.status.lock().expect("lock poisoned").record_error(&e);
                }
                self.state = Some(DataWriterState::Closed);
                Poll::Ready(Err(e.into()))
            }
            Poll::Ready((mut imp, Ok(()))) => {
                // Great!  We're done flushing.  Queue as much as we can of this
                // cell.
                let n_queued = imp.queue_bytes(buf);
                self.state = Some(DataWriterState::Ready(imp));
                Poll::Ready(Ok(n_queued))
            }
            Poll::Pending => {
                self.state = Some(DataWriterState::Flushing(future));
                Poll::Pending
            }
        }
    }

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

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

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

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

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

impl DataWriterImpl {
    /// Try to flush the current buffer contents as a data cell.
    async fn flush_buf(mut self) -> (Self, Result<()>) {
        let result =
            if let Some((cell, remainder)) = Data::try_split_from(&self.buf[..self.n_pending]) {
                // TODO: Eventually we may want a larger buffer; if we do,
                // this invariant will become false.
                assert!(remainder.is_empty());
                self.n_pending = 0;
                self.s.send(cell.into()).await
            } else {
                Ok(())
            };

        (self, result)
    }

    /// Add as many bytes as possible from `b` to our internal buffer;
    /// return the number we were able to add.
    fn queue_bytes(&mut self, b: &[u8]) -> usize {
        let empty_space = &mut self.buf[self.n_pending..];
        if empty_space.is_empty() {
            // that is, len == 0
            return 0;
        }

        let n_to_copy = std::cmp::min(b.len(), empty_space.len());
        empty_space[..n_to_copy].copy_from_slice(&b[..n_to_copy]);
        self.n_pending += n_to_copy;
        n_to_copy
    }
}

impl DataReader {
    /// Return a [`DataStreamCtrl`] object that can be used to monitor and
    /// interact with this stream without holding the stream itself.
    #[cfg(feature = "stream-ctrl")]
    pub fn ctrl(&self) -> &Arc<DataStreamCtrl> {
        &self.ctrl
    }
}

/// An enumeration for the state of a DataReader.
///
/// We have to use an enum here because, when we're waiting for
/// ReadingCell to complete, the future returned by `read_cell()` owns the
/// DataCellImpl.  If we wanted to store the future and the cell at the
/// same time, we'd need to make a self-referential structure, which isn't
/// possible in safe Rust AIUI.
#[derive(Educe)]
#[educe(Debug)]
enum DataReaderState {
    /// In this state we have received an end cell or an error.
    Closed,
    /// In this state the reader is not currently fetching a cell; it
    /// either has data or not.
    Ready(DataReaderImpl),
    /// The reader is currently fetching a cell: this future is the
    /// progress it is making.
    ReadingCell(
        #[educe(Debug(method = "skip_fmt"))]
        Pin<Box<dyn Future<Output = (DataReaderImpl, Result<()>)> + Send>>,
    ),
}

/// Wrapper for the read part of a DataStream
#[derive(Educe)]
#[educe(Debug)]
struct DataReaderImpl {
    /// The underlying StreamReader object.
    #[educe(Debug(method = "skip_fmt"))]
    s: StreamReader,

    /// 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 DataReader {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<IoResult<usize>> {
        // We're pulling the state object out of the reader.  We MUST
        // put it back before this function returns.
        let mut state = self.state.take().expect("Missing state in DataReader");

        loop {
            let mut future = match state {
                DataReaderState::Ready(mut imp) => {
                    // There may be data to read already.
                    let n_copied = imp.extract_bytes(buf);
                    if n_copied != 0 {
                        // We read data into the buffer.  Tell the caller.
                        self.state = Some(DataReaderState::Ready(imp));
                        return Poll::Ready(Ok(n_copied));
                    }

                    // No data available!  We have to launch a read.
                    Box::pin(imp.read_cell())
                }
                DataReaderState::ReadingCell(fut) => fut,
                DataReaderState::Closed => {
                    self.state = Some(DataReaderState::Closed);
                    return Poll::Ready(Err(Error::NotConnected.into()));
                }
            };

            // We have a future that represents an in-progress read.
            // See if it can make progress.
            match future.as_mut().poll(cx) {
                Poll::Ready((_imp, Err(e))) => {
                    // There aren't any survivable errors in the current
                    // design.
                    self.state = Some(DataReaderState::Closed);
                    #[cfg(feature = "stream-ctrl")]
                    {
                        _imp.status.lock().expect("lock poisoned").record_error(&e);
                    }
                    let result = if matches!(e, Error::EndReceived(EndReason::DONE)) {
                        Ok(0)
                    } else {
                        Err(e.into())
                    };
                    return Poll::Ready(result);
                }
                Poll::Ready((imp, Ok(()))) => {
                    // It read a cell!  Continue the loop.
                    state = DataReaderState::Ready(imp);
                }
                Poll::Pending => {
                    // The future is pending; store it and tell the
                    // caller to get back to us later.
                    self.state = Some(DataReaderState::ReadingCell(future));
                    return Poll::Pending;
                }
            }
        }
    }
}

#[cfg(feature = "tokio")]
impl TokioAsyncRead for DataReader {
    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.
    fn extract_bytes(&mut self, buf: &mut [u8]) -> usize {
        let remainder = &self.pending[self.offset..];
        let n_to_copy = std::cmp::min(buf.len(), remainder.len());
        buf[..n_to_copy].copy_from_slice(&remainder[..n_to_copy]);
        self.offset += n_to_copy;

        n_to_copy
    }

    /// Return true iff there are no buffered bytes here to yield
    fn buf_is_empty(&self) -> bool {
        self.pending.len() == self.offset
    }

    /// Load self.pending with the contents of a new data cell.
    ///
    /// This function takes ownership of self so that we can avoid
    /// self-referential lifetimes.
    async fn read_cell(mut self) -> (Self, Result<()>) {
        use DataStreamMsg::*;
        let msg = match self.s.recv().await {
            Ok(unparsed) => match unparsed.decode::<DataStreamMsg>() {
                Ok(cell) => cell.into_msg(),
                Err(e) => {
                    self.s.protocol_error();
                    return (
                        self,
                        Err(Error::from_bytes_err(e, "message on a data stream")),
                    );
                }
            },
            Err(e) => return (self, Err(e)),
        };

        let result = match msg {
            Connected(_) if !self.connected => {
                self.connected = true;
                #[cfg(feature = "stream-ctrl")]
                {
                    self.status
                        .lock()
                        .expect("poisoned lock")
                        .record_connected();
                }
                Ok(())
            }
            Connected(_) => {
                self.s.protocol_error();
                Err(Error::StreamProto(
                    "Received a second connect cell on a data stream".to_string(),
                ))
            }
            Data(d) if self.connected => {
                self.add_data(d.into());
                Ok(())
            }
            Data(_) => {
                self.s.protocol_error();
                Err(Error::StreamProto(
                    "Received a data cell an unconnected stream".to_string(),
                ))
            }
            End(e) => Err(Error::EndReceived(e.reason())),
        };

        (self, result)
    }

    /// Add the data from `d` to the end of our pending bytes.
    fn add_data(&mut self, mut d: Vec<u8>) {
        if self.buf_is_empty() {
            // No data pending?  Just take d as the new pending.
            self.pending = d;
            self.offset = 0;
        } 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);
        }
    }
}

/// A `CmdChecker` that enforces invariants for outbound data streams.
#[derive(Debug)]
pub(crate) struct DataCmdChecker {
    /// True if we are expecting to receive a CONNECTED message on this stream.
    expecting_connected: bool,
}

impl Default for DataCmdChecker {
    fn default() -> Self {
        Self {
            expecting_connected: true,
        }
    }
}

impl super::CmdChecker for DataCmdChecker {
    fn check_msg(
        &mut self,
        msg: &tor_cell::relaycell::UnparsedRelayCell,
    ) -> Result<super::StreamStatus> {
        use super::StreamStatus::*;
        match msg.cmd() {
            RelayCmd::CONNECTED => {
                if !self.expecting_connected {
                    Err(Error::StreamProto(
                        "Received CONNECTED twice on a stream.".into(),
                    ))
                } else {
                    self.expecting_connected = false;
                    Ok(Open)
                }
            }
            RelayCmd::DATA => {
                if !self.expecting_connected {
                    Ok(Open)
                } else {
                    Err(Error::StreamProto(
                        "Received DATA before CONNECTED on a stream".into(),
                    ))
                }
            }
            RelayCmd::END => Ok(Closed),
            _ => Err(Error::StreamProto(format!(
                "Unexpected {} on a data stream!",
                msg.cmd()
            ))),
        }
    }

    fn consume_checked_msg(&mut self, msg: tor_cell::relaycell::UnparsedRelayCell) -> Result<()> {
        let _ = msg
            .decode::<DataStreamMsg>()
            .map_err(|err| Error::from_bytes_err(err, "cell on half-closed stream"))?;
        Ok(())
    }
}

impl DataCmdChecker {
    /// Return a new boxed `DataCmdChecker` in a state suitable for a newly
    /// constructed connection.
    pub(crate) fn new_any() -> AnyCmdChecker {
        Box::<Self>::default()
    }

    /// Return a new boxed `DataCmdChecker` in a state suitable for a
    /// connection where an initial CONNECTED cell is not expected.
    ///
    /// This is used by hidden services, exit relays, and directory servers
    /// to accept streams.
    #[cfg(feature = "hs-service")]
    pub(crate) fn new_connected() -> AnyCmdChecker {
        Box::new(Self {
            expecting_connected: false,
        })
    }
}