impl RTCHandler for StunCodec

webrtc-rs · Mar 24, 2024 · a04c8ce · a04c8ce
1 parent cd660f2
commit a04c8ce
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 38 deletions.
diff --git a/rtc/src/handlers/mod.rs b/rtc/src/handlers/mod.rs
@@ -5,7 +5,7 @@ use std::time::Instant;
 
 pub mod demuxer;
 pub mod dtls;
-//pub mod stun;
+pub mod stun;
 
 pub trait RTCHandler {
     /// Handles input message

diff --git a/rtc/src/handlers/stun.rs b/rtc/src/handlers/stun.rs
@@ -1,5 +1,5 @@
 use crate::handlers::RTCHandler;
-use crate::messages::{RTCMessageEvent, STUNMessageEvent};
+use crate::messages::{RTCMessage, STUNMessage};
 use bytes::BytesMut;
 use log::{debug, warn};
 use shared::error::Result;
@@ -8,28 +8,17 @@ use stun::message::Message;
 
 /// StunHandler implements STUN Protocol handling
 #[derive(Default)]
-pub struct StunHandler {
-    next: Option<Box<dyn RTCHandler>>,
-}
+pub struct StunCodec;
 
-impl StunHandler {
+impl StunCodec {
     pub fn new() -> Self {
-        StunHandler::default()
+        StunCodec
     }
 }
 
-impl RTCHandler for StunHandler {
-    fn chain(mut self: Box<Self>, next: Box<dyn RTCHandler>) -> Box<dyn RTCHandler> {
-        self.next = Some(next);
-        self
-    }
-
-    fn next(&mut self) -> Option<&mut Box<dyn RTCHandler>> {
-        self.next.as_mut()
-    }
-
-    fn handle_transmit(&mut self, msg: Transmit<RTCMessageEvent>) {
-        let next_msg = if let RTCMessageEvent::Stun(STUNMessageEvent::Raw(message)) = msg.message {
+impl RTCHandler for StunCodec {
+    fn handle_transmit(&mut self, msg: Transmit<RTCMessage>) -> Vec<Transmit<RTCMessage>> {
+        if let RTCMessage::Stun(STUNMessage::Raw(message)) = msg.message {
             let try_read = || -> Result<Message> {
                 let mut stun_message = Message {
                     raw: message.to_vec(),
@@ -44,36 +33,26 @@ impl RTCHandler for StunHandler {
             };
 
             match try_read() {
-                Ok(stun_message) => Transmit {
+                Ok(stun_message) => vec![Transmit {
                     now: msg.now,
                     transport: msg.transport,
-                    message: RTCMessageEvent::Stun(STUNMessageEvent::Stun(stun_message)),
-                },
+                    message: RTCMessage::Stun(STUNMessage::Stun(stun_message)),
+                }],
                 Err(err) => {
                     warn!("try_read got error {}", err);
                     self.handle_error(err);
-                    return;
+                    vec![]
                 }
             }
         } else {
             debug!("bypass StunHandler read for {}", msg.transport.peer_addr);
-            msg
-        };
-
-        if let Some(next) = self.next() {
-            next.handle_transmit(next_msg);
+            vec![msg]
         }
     }
 
-    fn poll_transmit(&mut self) -> Option<Transmit<RTCMessageEvent>> {
-        let transmit = if let Some(next) = self.next() {
-            next.poll_transmit()
-        } else {
-            None
-        };
-
-        if let Some(msg) = transmit {
-            if let RTCMessageEvent::Stun(STUNMessageEvent::Stun(mut stun_message)) = msg.message {
+    fn poll_transmit(&mut self, msg: Option<Transmit<RTCMessage>>) -> Option<Transmit<RTCMessage>> {
+        if let Some(msg) = msg {
+            if let RTCMessage::Stun(STUNMessage::Stun(mut stun_message)) = msg.message {
                 debug!(
                     "StunMessage type {} sent to {}",
                     stun_message.typ, msg.transport.peer_addr
@@ -83,7 +62,7 @@ impl RTCHandler for StunHandler {
                 Some(Transmit {
                     now: msg.now,
                     transport: msg.transport,
-                    message: RTCMessageEvent::Stun(STUNMessageEvent::Raw(message)),
+                    message: RTCMessage::Stun(STUNMessage::Raw(message)),
                 })
             } else {
                 debug!("bypass StunHandler write for {}", msg.transport.peer_addr);