fix: multistream-select negotiation on outbound substream over webrtc

.gitignore

@@ -1 +1,3 @@
 /target
+.idea
+

Cargo.toml

@@ -65,7 +65,7 @@ rcgen = { version = "0.14.5", optional = true }
 # End of Quic related dependencies.
 
 # WebRTC related dependencies. WebRTC is an experimental feature flag. The dependencies must be updated.
-str0m = { version = "0.9.0", optional = true }
+str0m = { version = "0.11.1", optional = true }
 # End of WebRTC related dependencies.
 
 # Fuzzing related dependencies.

src/multistream_select/dialer_select.rs

@@ -22,11 +22,11 @@
 
 use crate::{
     codec::unsigned_varint::UnsignedVarint,
-    error::{self, Error, ParseError},
+    error::{self, Error, ParseError, SubstreamError},
     multistream_select::{
         protocol::{
             webrtc_encode_multistream_message, HeaderLine, Message, MessageIO, Protocol,
-            ProtocolError,
+            ProtocolError, PROTO_MULTISTREAM_1_0
         },
         Negotiated, NegotiationError, Version,
     },
@@ -357,24 +357,95 @@ impl WebRtcDialerState {
         &mut self,
         payload: Vec<u8>,
     ) -> Result<HandshakeResult, crate::error::NegotiationError> {
-        let Message::Protocols(protocols) =
-            Message::decode(payload.into()).map_err(|_| ParseError::InvalidData)?
-        else {
-            return Err(crate::error::NegotiationError::MultistreamSelectError(
-                NegotiationError::Failed,
-            ));
+        // All multistream-select messages are length-prefixed. Since this code path is not using
+        // multistream_select::protocol::MessageIO, we need to decode and remove the length here.
+        let mut remaining: &[u8] = &payload;
+        let (len, tail) = unsigned_varint::decode::usize(remaining).
+            map_err(|error| {
+                tracing::debug!(
+                    target: LOG_TARGET,
+                    ?error,
+                    message = ?payload,
+                    "Failed to decode length-prefix in multistream message");
+                error::NegotiationError::ParseError(ParseError::InvalidData)
+            })?;
+
+        let payload = tail[..len].to_vec();
+
+        let message = Message::decode(payload.into());
+
+        tracing::trace!(
+            target: LOG_TARGET,
+            ?message,
+            "Decoded message while registering response",
+        );
+
+        let mut protocols = match message {
+            Ok(Message::Header(HeaderLine::V1)) => {
+                vec![PROTO_MULTISTREAM_1_0]
+            }
+            Ok(Message::Protocol(protocol)) => vec![protocol],
+            Ok(Message::Protocols(protocols)) => protocols,
+            Ok(Message::NotAvailable) => {
+                return match &self.state {
+                    HandshakeState::WaitingProtocol =>
+                        Err(error::NegotiationError::MultistreamSelectError(
+                            NegotiationError::Failed
+                        )),
+                    _ => Err(error::NegotiationError::StateMismatch),
+                }
+            }
+            Ok(Message::ListProtocols) => return Err(error::NegotiationError::StateMismatch),
+            Err(_) => return Err(error::NegotiationError::ParseError(ParseError::InvalidData)),
         };
 
+        remaining = &tail[len..];
+
+        loop {
+            if remaining.is_empty() {
+                break;
+            }
+
+            let (len, tail) = unsigned_varint::decode::usize(remaining).
+                map_err(|error| {
+                    tracing::debug!(
+                    target: LOG_TARGET,
+                    ?error,
+                    message = ?remaining,
+                    "Failed to decode length-prefix in multistream message");
+                    error::NegotiationError::ParseError(ParseError::InvalidData)
+                })?;
+
+            if len > tail.len() {
+                break;
+            }
+
+            let payload = tail[..len].to_vec();
+
+            match Message::decode(payload.into()) {
+                Ok(Message::Protocol(protocol)) => protocols.push(protocol),
+                Err(error) => {
+                    tracing::debug!(
+                        target: LOG_TARGET,
+                        ?error,
+                        message = ?tail[..len],
+                        "Failed to decode multistream message",
+                    );
+                    return Err(error::NegotiationError::ParseError(ParseError::InvalidData))
+                },
+                _ => return Err(error::NegotiationError::StateMismatch),
+            }
+
+            remaining = &tail[len..];
+        }
+
         let mut protocol_iter = protocols.into_iter();
         loop {
             match (&self.state, protocol_iter.next()) {
                 (HandshakeState::WaitingResponse, None) =>
                     return Err(crate::error::NegotiationError::StateMismatch),
                 (HandshakeState::WaitingResponse, Some(protocol)) => {
-                    let header = Protocol::try_from(&b"/multistream/1.0.0"[..])
-                        .expect("valid multitstream-select header");
-
-                    if protocol == header {
+                    if protocol == PROTO_MULTISTREAM_1_0 {
                         self.state = HandshakeState::WaitingProtocol;
                     } else {
                         return Err(crate::error::NegotiationError::MultistreamSelectError(
@@ -410,7 +481,9 @@ mod tests {
     use super::*;
     use crate::multistream_select::listener_select_proto;
     use std::time::Duration;
+    use bytes::BufMut;
     use tokio::net::{TcpListener, TcpStream};
+    use crate::multistream_select::protocol::MSG_MULTISTREAM_1_0;
 
     #[tokio::test]
     async fn select_proto_basic() {
@@ -755,23 +828,18 @@ mod tests {
     fn propose() {
         let (mut dialer_state, message) =
             WebRtcDialerState::propose(ProtocolName::from("/13371338/proto/1"), vec![]).unwrap();
-        let message = bytes::BytesMut::from(&message[..]).freeze();
 
-        let Message::Protocols(protocols) = Message::decode(message).unwrap() else {
-            panic!("invalid message type");
-        };
+        let mut bytes = BytesMut::with_capacity(32);
+        bytes.put_u8(MSG_MULTISTREAM_1_0.len() as u8);
+        let _ = Message::Header(HeaderLine::V1).encode(&mut bytes).unwrap();
 
-        assert_eq!(protocols.len(), 2);
-        assert_eq!(
-            protocols[0],
-            Protocol::try_from(&b"/multistream/1.0.0"[..])
-                .expect("valid multitstream-select header")
-        );
-        assert_eq!(
-            protocols[1],
-            Protocol::try_from(&b"/13371338/proto/1"[..])
-                .expect("valid multitstream-select header")
-        );
+        let proto = Protocol::try_from(&b"/13371338/proto/1"[..]).expect("valid protocol name");
+        bytes.put_u8((proto.as_ref().len() + 1) as u8); // + 1 for \n
+        let _ = Message::Protocol(proto).encode(&mut bytes).unwrap();
+
+        let expected_message = bytes.freeze().to_vec();
+
+        assert_eq!(message, expected_message);
     }
 
     #[test]
@@ -781,41 +849,38 @@ mod tests {
             vec![ProtocolName::from("/sup/proto/1")],
         )
         .unwrap();
-        let message = bytes::BytesMut::from(&message[..]).freeze();
 
-        let Message::Protocols(protocols) = Message::decode(message).unwrap() else {
-            panic!("invalid message type");
-        };
+        let mut bytes = BytesMut::with_capacity(32);
+        bytes.put_u8(MSG_MULTISTREAM_1_0.len() as u8);
+        let _ = Message::Header(HeaderLine::V1).encode(&mut bytes).unwrap();
 
-        assert_eq!(protocols.len(), 3);
-        assert_eq!(
-            protocols[0],
-            Protocol::try_from(&b"/multistream/1.0.0"[..])
-                .expect("valid multitstream-select header")
-        );
-        assert_eq!(
-            protocols[1],
-            Protocol::try_from(&b"/13371338/proto/1"[..])
-                .expect("valid multitstream-select header")
-        );
-        assert_eq!(
-            protocols[2],
-            Protocol::try_from(&b"/sup/proto/1"[..]).expect("valid multitstream-select header")
-        );
+        let proto1 = Protocol::try_from(&b"/13371338/proto/1"[..]).expect("valid protocol name");
+        bytes.put_u8((proto1.as_ref().len() + 1) as u8); // + 1 for \n
+        let _ = Message::Protocol(proto1).encode(&mut bytes).unwrap();
+
+        let proto2 = Protocol::try_from(&b"/sup/proto/1"[..]).expect("valid protocol name");
+        bytes.put_u8((proto2.as_ref().len() + 1) as u8); // + 1 for \n
+        let _ = Message::Protocol(proto2).encode(&mut bytes).unwrap();
+
+        let expected_message = bytes.freeze().to_vec();
+
+        assert_eq!(message, expected_message);
     }
 
     #[test]
-    fn register_response_invalid_message() {
-        // send only header line
+    fn register_response_header_only() {
         let mut bytes = BytesMut::with_capacity(32);
+        bytes.put_u8(MSG_MULTISTREAM_1_0.len() as u8);
+
         let message = Message::Header(HeaderLine::V1);
         message.encode(&mut bytes).map_err(|_| Error::InvalidData).unwrap();
 
         let (mut dialer_state, _message) =
             WebRtcDialerState::propose(ProtocolName::from("/13371338/proto/1"), vec![]).unwrap();
 
         match dialer_state.register_response(bytes.freeze().to_vec()) {
-            Err(error::NegotiationError::MultistreamSelectError(NegotiationError::Failed)) => {}
+            Ok(HandshakeResult::NotReady) => {},
+            Err(err) => panic!("unexpected error: {:?}", err),
             event => panic!("invalid event: {event:?}"),
         }
     }

src/multistream_select/listener_select.rs

@@ -351,6 +351,14 @@ pub fn webrtc_listener_negotiate<'a>(
     supported_protocols: &'a mut impl Iterator<Item = &'a ProtocolName>,
     payload: Bytes,
 ) -> crate::Result<ListenerSelectResult> {
+    let payload = if payload.len() > 2 && payload[0..payload.len() - 2] != b"\n\n"[..] {
+        let mut buf = BytesMut::from(payload);
+        buf.extend_from_slice(b"\n");
+        buf.freeze()
+    } else {
+        payload
+    };
+
     let Message::Protocols(protocols) = Message::decode(payload).map_err(|_| Error::InvalidData)?
     else {
         return Err(Error::NegotiationError(
@@ -466,12 +474,7 @@ mod tests {
         message.encode(&mut bytes).map_err(|_| Error::InvalidData).unwrap();
 
         match webrtc_listener_negotiate(&mut local_protocols.iter(), bytes.freeze()) {
-            Err(error) => assert!(std::matches!(
-                error,
-                Error::NegotiationError(error::NegotiationError::MultistreamSelectError(
-                    NegotiationError::Failed
-                ))
-            )),
+            Err(error) => assert!(std::matches!(error, Error::InvalidData)),
             event => panic!("invalid event: {event:?}"),
         }
     }
@@ -495,12 +498,7 @@ mod tests {
         message.encode(&mut bytes).map_err(|_| Error::InvalidData).unwrap();
 
         match webrtc_listener_negotiate(&mut local_protocols.iter(), bytes.freeze()) {
-            Err(error) => assert!(std::matches!(
-                error,
-                Error::NegotiationError(error::NegotiationError::MultistreamSelectError(
-                    NegotiationError::Failed
-                ))
-            )),
+            Err(error) => assert!(std::matches!(error, Error::InvalidData)),
             event => panic!("invalid event: {event:?}"),
         }
     }

src/multistream_select/protocol.rs

@@ -54,6 +54,8 @@ pub const MSG_MULTISTREAM_1_0: &[u8] = b"/multistream/1.0.0\n";
 const MSG_PROTOCOL_NA: &[u8] = b"na\n";
 /// The encoded form of a multistream-select 'ls' message.
 const MSG_LS: &[u8] = b"ls\n";
+/// A Protocol instance for the `/multistream/1.0.0` header line.
+pub const PROTO_MULTISTREAM_1_0: Protocol = Protocol(Bytes::from_static(MSG_MULTISTREAM_1_0));
 /// Logging target.
 const LOG_TARGET: &str = "litep2p::multistream-select";
 
@@ -230,7 +232,7 @@ impl Message {
 ///
 /// # Note
 ///
-/// This is implementation is not compliant with the multistream-select protocol spec.
+/// This implementation may not be compliant with the multistream-select protocol spec.
 /// The only purpose of this was to get the `multistream-select` protocol working with smoldot.
 pub fn webrtc_encode_multistream_message(
     messages: impl IntoIterator<Item = Message>,
@@ -249,9 +251,6 @@ pub fn webrtc_encode_multistream_message(
         header.append(&mut proto_bytes);
     }
 
-    // For the `Message::Protocols` to be interpreted correctly, it must be followed by a newline.
-    header.push(b'\n');
-
     Ok(BytesMut::from(&header[..]))
 }