feat(webrtc): implement fin ack

misc/webrtc-utils/CHANGELOG.md

@@ -1,3 +1,8 @@
+## 0.3.1
+
+- Implement FIN_ACK support for WebRTC streams.
+  See [PR 5687](https://github.com/libp2p/rust-libp2p/pull/5687).
+
 ## 0.3.0
 
 <!-- Update to libp2p-swarm v0.45.0 -->
@@ -15,6 +20,6 @@
 ## 0.1.0
 
 - Initial release.
-  See [PR 4248].
+  See [PR 4248](https://github.com/libp2p/rust-libp2p/pull/4248).
 
 [PR 4248]: https://github.com/libp2p/rust-libp2p/pull/4248

misc/webrtc-utils/src/generated/message.proto

@@ -12,6 +12,10 @@ message Message {
     // The sender abruptly terminates the sending part of the stream. The
     // receiver can discard any data that it already received on that stream.
     RESET = 2;
+    // Sending the FIN_ACK flag acknowledges the previous receipt of a message
+    // with the FIN flag set. Receiving a FIN_ACK flag gives the recipient
+    // confidence that the remote has received all sent messages.
+    FIN_ACK = 3;
   }
 
   optional Flag flag=1;

misc/webrtc-utils/src/generated/webrtc/pb.rs

@@ -57,6 +57,7 @@ pub enum Flag {
     FIN = 0,
     STOP_SENDING = 1,
     RESET = 2,
+    FIN_ACK = 3,
 }
 
 impl Default for Flag {
@@ -71,6 +72,7 @@ impl From<i32> for Flag {
             0 => Flag::FIN,
             1 => Flag::STOP_SENDING,
             2 => Flag::RESET,
+            3 => Flag::FIN_ACK,
             _ => Self::default(),
         }
     }
@@ -82,6 +84,7 @@ impl<'a> From<&'a str> for Flag {
             "FIN" => Flag::FIN,
             "STOP_SENDING" => Flag::STOP_SENDING,
             "RESET" => Flag::RESET,
+            "FIN_ACK" => Flag::FIN_ACK,
             _ => Self::default(),
         }
     }

misc/webrtc-utils/src/stream.rs

@@ -23,6 +23,7 @@ use std::{
     io,
     pin::Pin,
     task::{Context, Poll},
+    time::{Duration, Instant},
 };
 
 use bytes::Bytes;
@@ -53,6 +54,8 @@ const VARINT_LEN: usize = 2;
 const PROTO_OVERHEAD: usize = 5;
 /// Maximum length of data, in bytes.
 const MAX_DATA_LEN: usize = MAX_MSG_LEN - VARINT_LEN - PROTO_OVERHEAD;
+/// FIN_ACK timeout
+const FIN_ACK_TIMEOUT: Duration = Duration::from_secs(10);
 
 pub use drop_listener::DropListener;
 /// A stream backed by a WebRTC data channel.
@@ -65,6 +68,7 @@ pub struct Stream<T> {
     read_buffer: Bytes,
     /// Dropping this will close the oneshot and notify the receiver by emitting `Canceled`.
     drop_notifier: Option<oneshot::Sender<GracefullyClosed>>,
+    fin_ack_deadline: Option<Instant>,
 }
 
 impl<T> Stream<T>
@@ -81,6 +85,7 @@ where
             state: State::Open,
             read_buffer: Bytes::default(),
             drop_notifier: Some(sender),
+            fin_ack_deadline: None,
         };
         let listener = DropListener::new(framed_dc::new(data_channel), receiver);
 
@@ -146,6 +151,16 @@ where
                 Some((flag, message)) => {
                     if let Some(flag) = flag {
                         state.handle_inbound_flag(flag, read_buffer);
+
+                        // Send FIN_ACK in response to FIN
+                        if flag == Flag::FIN && state.should_send_fin_ack() {
+                            ready!(io.poll_ready_unpin(cx))?;
+                            io.start_send_unpin(Message {
+                                flag: Some(Flag::FIN_ACK),
+                                message: None,
+                            })?;
+                            ready!(io.poll_flush_unpin(cx))?;
+                        }
                     }
 
                     debug_assert!(read_buffer.is_empty());
@@ -231,19 +246,39 @@ where
                     })?;
                     self.state.close_write_message_sent();
 
+                    // Set deadline when sending FIN
+                    self.fin_ack_deadline = Some(Instant::now() + FIN_ACK_TIMEOUT);
+
                     continue;
                 }
                 Some(Closing::MessageSent) => {
                     ready!(self.io.poll_flush_unpin(cx))?;
 
-                    self.state.write_closed();
-                    let _ = self
-                        .drop_notifier
-                        .take()
-                        .expect("to not close twice")
-                        .send(GracefullyClosed {});
+                    if self.state.fin_ack_received() {
+                        self.state.write_closed();
+                        let _ = self
+                            .drop_notifier
+                            .take()
+                            .expect("to not close twice")
+                            .send(GracefullyClosed {});
+                        return Poll::Ready(Ok(()));
+                    }
 
-                    return Poll::Ready(Ok(()));
+                    if self
+                        .fin_ack_deadline
+                        .is_some_and(|deadline| Instant::now() >= deadline)
+                    {
+                        tracing::warn!("FIN_ACK timeout, forcing close");
+                        self.state.write_closed();
+                        let _ = self
+                            .drop_notifier
+                            .take()
+                            .expect("to not close twice")
+                            .send(GracefullyClosed {});
+                        return Poll::Ready(Ok(()));
+                    }
+
+                    return Poll::Pending;
                 }
                 None => return Poll::Ready(Ok(())),
             }

misc/webrtc-utils/src/stream/state.rs

@@ -38,6 +38,7 @@ pub(crate) enum State {
         /// Whether the write side of our channel was already closed.
         read_closed: bool,
         inner: Closing,
+        fin_ack_received: bool,
     },
     BothClosed {
         reset: bool,
@@ -72,6 +73,20 @@ impl State {
             (Self::ReadClosed, Flag::STOP_SENDING) => {
                 *self = Self::BothClosed { reset: false };
             }
+            (
+                Self::ClosingWrite {
+                    inner: Closing::MessageSent,
+                    read_closed,
+                    fin_ack_received: false,
+                },
+                Flag::FIN_ACK,
+            ) => {
+                *self = Self::ClosingWrite {
+                    read_closed,
+                    inner: Closing::MessageSent,
+                    fin_ack_received: true,
+                };
+            }
             (_, Flag::RESET) => {
                 buffer.clear();
                 *self = Self::BothClosed { reset: true };
@@ -85,6 +100,7 @@ impl State {
             State::ClosingWrite {
                 read_closed: true,
                 inner,
+                ..
             } => {
                 debug_assert!(matches!(inner, Closing::MessageSent));
 
@@ -93,6 +109,7 @@ impl State {
             State::ClosingWrite {
                 read_closed: false,
                 inner,
+                ..
             } => {
                 debug_assert!(matches!(inner, Closing::MessageSent));
 
@@ -110,12 +127,15 @@ impl State {
 
     pub(crate) fn close_write_message_sent(&mut self) {
         match self {
-            State::ClosingWrite { inner, read_closed } => {
+            State::ClosingWrite {
+                inner, read_closed, ..
+            } => {
                 debug_assert!(matches!(inner, Closing::Requested));
 
                 *self = State::ClosingWrite {
                     read_closed: *read_closed,
                     inner: Closing::MessageSent,
+                    fin_ack_received: false,
                 };
             }
             State::Open
@@ -244,12 +264,14 @@ impl State {
                     *self = Self::ClosingWrite {
                         read_closed: false,
                         inner: Closing::Requested,
+                        fin_ack_received: false,
                     };
                 }
                 State::ReadClosed => {
                     *self = Self::ClosingWrite {
                         read_closed: true,
                         inner: Closing::Requested,
+                        fin_ack_received: false,
                     };
                 }
 
@@ -320,6 +342,26 @@ impl State {
             }
         }
     }
+
+    /// Returns true if this state should acknowledge a received FIN flag with a FIN_ACK response, which occurs in ReadClosed, ClosingRead, or non-reset BothClosed states.
+    pub(crate) fn should_send_fin_ack(&self) -> bool {
+        matches!(
+            self,
+            Self::ReadClosed | Self::ClosingRead { .. } | Self::BothClosed { reset: false }
+        )
+    }
+
+    /// Returns true if we've received a FIN_ACK response in ClosingWrite state after sending our FIN message, indicating we can complete the write closure process.
+    pub(crate) fn fin_ack_received(&self) -> bool {
+        matches!(
+            self,
+            Self::ClosingWrite {
+                fin_ack_received: true,
+                inner: Closing::MessageSent,
+                ..
+            }
+        )
+    }
 }
 
 #[cfg(test)]
@@ -506,4 +548,130 @@ mod tests {
 
         assert!(buffer.is_empty());
     }
+
+    #[test]
+    fn should_send_fin_ack_when_read_closed() {
+        let mut state = State::Open;
+        state.handle_inbound_flag(Flag::FIN, &mut Bytes::default());
+        assert!(state.should_send_fin_ack());
+    }
+
+    #[test]
+    fn should_send_fin_ack_when_closing_read() {
+        let mut state = State::Open;
+        state.close_read_barrier().unwrap();
+        assert!(state.should_send_fin_ack());
+    }
+
+    #[test]
+    fn should_not_send_fin_ack_when_write_closed() {
+        let mut state = State::Open;
+        state.close_write_barrier().unwrap();
+        state.close_write_message_sent();
+        state.write_closed();
+        assert!(!state.should_send_fin_ack());
+    }
+
+    #[test]
+    fn fin_ack_received_transitions_to_write_closed() {
+        let mut state = State::Open;
+
+        // Start closing write side
+        state.close_write_barrier().unwrap();
+        state.close_write_message_sent();
+
+        // Receive FIN_ACK
+        state.handle_inbound_flag(Flag::FIN_ACK, &mut Bytes::default());
+        assert!(state.fin_ack_received());
+
+        // Call write_closed to complete transition
+        state.write_closed();
+
+        // Should transition to WriteClosed
+        match state {
+            State::WriteClosed => {}
+            _ => panic!("Expected WriteClosed state, got {:?}", state),
+        }
+    }
+
+    #[test]
+    fn fin_ack_received_transitions_to_both_closed() {
+        let mut state = State::Open;
+
+        // Close read side first
+        state.handle_inbound_flag(Flag::FIN, &mut Bytes::default());
+        assert!(matches!(state, State::ReadClosed));
+
+        // Start closing write side
+        state.close_write_barrier().unwrap();
+        state.close_write_message_sent();
+
+        // Receive FIN_ACK
+        state.handle_inbound_flag(Flag::FIN_ACK, &mut Bytes::default());
+        assert!(state.fin_ack_received());
+
+        // Call write_closed to complete transition
+        state.write_closed();
+
+        // Should transition to BothClosed
+        match state {
+            State::BothClosed { reset: false } => {}
+            _ => panic!("Expected BothClosed state, got {:?}", state),
+        }
+    }
+
+    #[test]
+    fn fin_ack_received_flag_is_set() {
+        let mut state = State::Open;
+
+        // Start closing write side
+        state.close_write_barrier().unwrap();
+        assert!(matches!(
+            state,
+            State::ClosingWrite {
+                fin_ack_received: false,
+                inner: Closing::Requested,
+                ..
+            }
+        ));
+
+        state.close_write_message_sent();
+        assert!(matches!(
+            state,
+            State::ClosingWrite {
+                fin_ack_received: false,
+                inner: Closing::MessageSent,
+                ..
+            }
+        ));
+
+        assert!(!state.fin_ack_received());
+
+        // Receive FIN_ACK
+        state.handle_inbound_flag(Flag::FIN_ACK, &mut Bytes::default());
+
+        assert!(state.fin_ack_received());
+        assert!(matches!(
+            state,
+            State::ClosingWrite {
+                fin_ack_received: true,
+                inner: Closing::MessageSent,
+                ..
+            }
+        ));
+    }
+
+    #[test]
+    fn ignore_fin_ack_in_wrong_state() {
+        let mut state = State::Open;
+
+        // Should ignore FIN_ACK in Open state
+        state.handle_inbound_flag(Flag::FIN_ACK, &mut Bytes::default());
+        assert!(matches!(state, State::Open));
+
+        // Should ignore FIN_ACK in ReadClosed state
+        state.handle_inbound_flag(Flag::FIN, &mut Bytes::default());
+        state.handle_inbound_flag(Flag::FIN_ACK, &mut Bytes::default());
+        assert!(matches!(state, State::ReadClosed));
+    }
 }