3 files changed, 155 insertions, 0 deletions
diff --git a/crates/tor-proto/src/stream/cmdcheck.rs b/crates/tor-proto/src/stream/cmdcheck.rs
new file mode 100644
index 000000000..76808e0df
--- /dev/null
+++ b/crates/tor-proto/src/stream/cmdcheck.rs
@@ -0,0 +1,55 @@
+//! Declare a "command checker" trait that checks whether a given relay message
+//! is acceptable on a given stream.
+
+use tor_cell::relaycell::UnparsedRelayCell;
+
+use crate::Result;
+
+/// A value returned by CmdChecker on success.
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+pub(crate) enum StreamStatus {
+    /// The stream is still open.
+    Open,
+    /// The stream has been closed successfully; any further messages received
+    /// on this stream would be a protocol violation, which should cause
+    /// us to close the circuit.
+    Closed,
+}
+
+/// An object that checks incoming commands before they are sent to a stream.
+///
+/// These checks are called from the circuit reactor code, which runs in its own
+/// task. The reactor code continues calling these checks we have sent our own
+/// END cell on the stream.  See `crate::circuit::halfstream` for more
+/// information.
+///
+/// NOTE: The checking DOES NOT take SENDME messages into account; those are
+/// handled separately.  Neither of the methods on this trait will ever be
+/// passed a SENDME message.
+///
+/// See [`circuit::reactor`](crate::circuit::reactor) for more information on
+/// how these checks relate to other checks performed on incoming messages.
+pub(crate) trait CmdChecker: std::fmt::Debug {
+    /// Look at a message `msg` and decide whether it can be handled on this
+    /// stream.
+    ///
+    /// If `msg` is invalid, return an error, indicating that the protocol has
+    /// been violated and the corresponding circuit should be closed.
+    ///
+    /// If `msg` is invalid, update the state of this checker, and return a
+    /// `StreamStatus` indicating whether the last message closed.
+    fn check_msg(&mut self, msg: &UnparsedRelayCell) -> Result<StreamStatus>;
+
+    /// Consume `msg` and make sure it can be parsed correctly.
+    ///
+    /// This is an additional check, beyond check_msg(), performed for half-open
+    /// streams.  It should only be called  if check_msg() succeeds.  It shouldn't
+    /// be called on open streams: for those, the stream itself parses the message
+    /// and consumes it.
+    fn consume_checked_msg(&mut self, msg: UnparsedRelayCell) -> Result<()>;
+}
+
+/// Type alias for a CmdChecker of unspecified type.
+//
+// TODO: Someday we might turn this into an enum if we decide it's beneficial.
+pub(crate) type AnyCmdChecker = Box<dyn CmdChecker + Send + 'static>;
diff --git a/crates/tor-proto/src/stream/data.rs b/crates/tor-proto/src/stream/data.rs
index 1d460e801..24a097713 100644
--- a/crates/tor-proto/src/stream/data.rs
+++ b/crates/tor-proto/src/stream/data.rs
@@ -3,6 +3,7 @@
 
 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};
@@ -28,6 +29,8 @@ 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
@@ -657,3 +660,60 @@ impl DataReaderImpl {
         }
     }
 }
+
+/// A `CmdChecker` that enforces invariants for outbound data streams.
+#[derive(Debug, Default)]
+pub(crate) struct DataCmdChecker {
+    /// True if we have received a CONNECTED message on this stream.
+    connected_received: bool,
+}
+
+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.connected_received {
+                    Err(Error::StreamProto(
+                        "Received CONNECTED twice on a stream.".into(),
+                    ))
+                } else {
+                    self.connected_received = true;
+                    Ok(Open)
+                }
+            }
+            RelayCmd::DATA => {
+                if self.connected_received {
+                    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()
+    }
+}
diff --git a/crates/tor-proto/src/stream/resolve.rs b/crates/tor-proto/src/stream/resolve.rs
index e81823e3d..f5c84b155 100644
--- a/crates/tor-proto/src/stream/resolve.rs
+++ b/crates/tor-proto/src/stream/resolve.rs
@@ -3,8 +3,11 @@
 use crate::stream::StreamReader;
 use crate::{Error, Result};
 use tor_cell::relaycell::msg::Resolved;
+use tor_cell::relaycell::RelayCmd;
 use tor_cell::restricted_msg;
 
+use super::AnyCmdChecker;
+
 /// A ResolveStream represents a pending DNS request made with a RESOLVE
 /// cell.
 pub struct ResolveStream {
@@ -46,3 +49,40 @@ impl ResolveStream {
         }
     }
 }
+
+/// A `CmdChecker` that enforces correctness for incoming commands on an
+/// outbound resolve stream.
+#[derive(Debug, Default)]
+pub(crate) struct ResolveCmdChecker {}
+
+impl super::CmdChecker for ResolveCmdChecker {
+    fn check_msg(
+        &mut self,
+        msg: &tor_cell::relaycell::UnparsedRelayCell,
+    ) -> Result<super::StreamStatus> {
+        use super::StreamStatus::Closed;
+        match msg.cmd() {
+            RelayCmd::RESOLVED => Ok(Closed),
+            RelayCmd::END => Ok(Closed),
+            _ => Err(Error::StreamProto(format!(
+                "Unexpected {} on resolve stream",
+                msg.cmd()
+            ))),
+        }
+    }
+
+    fn consume_checked_msg(&mut self, msg: tor_cell::relaycell::UnparsedRelayCell) -> Result<()> {
+        let _ = msg
+            .decode::<ResolveResponseMsg>()
+            .map_err(|err| Error::from_bytes_err(err, "message on resolve stream."))?;
+        Ok(())
+    }
+}
+
+impl ResolveCmdChecker {
+    /// Return a new boxed `DataCmdChecker` in a state suitable for a newly
+    /// constructed connection.
+    pub(crate) fn new_any() -> AnyCmdChecker {
+        Box::<Self>::default()
+    }
+}