1 use anyhow::Result;
2 use clap::{AppSettings, Arg, Command};
3 use std::collections::HashMap;
4 use std::io::Write;
5 use std::sync::Arc;
6 use tokio::net::UdpSocket;
7 use tokio::time::Duration;
8 use webrtc::api::interceptor_registry::register_default_interceptors;
9 use webrtc::api::media_engine::{MediaEngine, MIME_TYPE_OPUS, MIME_TYPE_VP8};
10 use webrtc::api::APIBuilder;
11 use webrtc::ice_transport::ice_connection_state::RTCIceConnectionState;
12 use webrtc::ice_transport::ice_server::RTCIceServer;
13 use webrtc::interceptor::registry::Registry;
14 use webrtc::peer_connection::configuration::RTCConfiguration;
15 use webrtc::peer_connection::peer_connection_state::RTCPeerConnectionState;
16 use webrtc::peer_connection::sdp::session_description::RTCSessionDescription;
17 use webrtc::rtcp::payload_feedbacks::picture_loss_indication::PictureLossIndication;
18 use webrtc::rtp_transceiver::rtp_codec::{
19     RTCRtpCodecCapability, RTCRtpCodecParameters, RTPCodecType,
20 };
21 use webrtc::util::{Conn, Marshal, Unmarshal};
22 
23 #[derive(Clone)]
24 struct UdpConn {
25     conn: Arc<dyn Conn + Send + Sync>,
26     payload_type: u8,
27 }
28 
29 #[tokio::main]
main() -> Result<()>30 async fn main() -> Result<()> {
31     let mut app = Command::new("rtp-forwarder")
32         .version("0.1.0")
33         .author("Rain Liu <[email protected]>")
34         .about("An example of rtp-forwarder.")
35         .setting(AppSettings::DeriveDisplayOrder)
36         .subcommand_negates_reqs(true)
37         .arg(
38             Arg::new("FULLHELP")
39                 .help("Prints more detailed help information")
40                 .long("fullhelp"),
41         )
42         .arg(
43             Arg::new("debug")
44                 .long("debug")
45                 .short('d')
46                 .help("Prints debug log information"),
47         );
48 
49     let matches = app.clone().get_matches();
50 
51     if matches.is_present("FULLHELP") {
52         app.print_long_help().unwrap();
53         std::process::exit(0);
54     }
55 
56     let debug = matches.is_present("debug");
57     if debug {
58         env_logger::Builder::new()
59             .format(|buf, record| {
60                 writeln!(
61                     buf,
62                     "{}:{} [{}] {} - {}",
63                     record.file().unwrap_or("unknown"),
64                     record.line().unwrap_or(0),
65                     record.level(),
66                     chrono::Local::now().format("%H:%M:%S.%6f"),
67                     record.args()
68                 )
69             })
70             .filter(None, log::LevelFilter::Trace)
71             .init();
72     }
73 
74     // Everything below is the WebRTC-rs API! Thanks for using it ❤️.
75 
76     // Create a MediaEngine object to configure the supported codec
77     let mut m = MediaEngine::default();
78 
79     // Setup the codecs you want to use.
80     // We'll use a VP8 and Opus but you can also define your own
81     m.register_codec(
82         RTCRtpCodecParameters {
83             capability: RTCRtpCodecCapability {
84                 mime_type: MIME_TYPE_VP8.to_owned(),
85                 clock_rate: 90000,
86                 channels: 0,
87                 sdp_fmtp_line: "".to_owned(),
88                 rtcp_feedback: vec![],
89             },
90             payload_type: 96,
91             ..Default::default()
92         },
93         RTPCodecType::Video,
94     )?;
95 
96     m.register_codec(
97         RTCRtpCodecParameters {
98             capability: RTCRtpCodecCapability {
99                 mime_type: MIME_TYPE_OPUS.to_owned(),
100                 clock_rate: 48000,
101                 channels: 2,
102                 sdp_fmtp_line: "".to_owned(),
103                 rtcp_feedback: vec![],
104             },
105             payload_type: 111,
106             ..Default::default()
107         },
108         RTPCodecType::Audio,
109     )?;
110 
111     // Create a InterceptorRegistry. This is the user configurable RTP/RTCP Pipeline.
112     // This provides NACKs, RTCP Reports and other features. If you use `webrtc.NewPeerConnection`
113     // this is enabled by default. If you are manually managing You MUST create a InterceptorRegistry
114     // for each PeerConnection.
115     let mut registry = Registry::new();
116 
117     // Use the default set of Interceptors
118     registry = register_default_interceptors(registry, &mut m)?;
119 
120     // Create the API object with the MediaEngine
121     let api = APIBuilder::new()
122         .with_media_engine(m)
123         .with_interceptor_registry(registry)
124         .build();
125 
126     // Prepare the configuration
127     let config = RTCConfiguration {
128         ice_servers: vec![RTCIceServer {
129             urls: vec!["stun:stun.l.google.com:19302".to_owned()],
130             ..Default::default()
131         }],
132         ..Default::default()
133     };
134 
135     // Create a new RTCPeerConnection
136     let peer_connection = Arc::new(api.new_peer_connection(config).await?);
137 
138     // Allow us to receive 1 audio track, and 1 video track
139     peer_connection
140         .add_transceiver_from_kind(RTPCodecType::Audio, None)
141         .await?;
142     peer_connection
143         .add_transceiver_from_kind(RTPCodecType::Video, None)
144         .await?;
145 
146     // Prepare udp conns
147     // Also update incoming packets with expected PayloadType, the browser may use
148     // a different value. We have to modify so our stream matches what rtp-forwarder.sdp expects
149     let mut udp_conns = HashMap::new();
150     udp_conns.insert(
151         "audio".to_owned(),
152         UdpConn {
153             conn: {
154                 let sock = UdpSocket::bind("127.0.0.1:0").await?;
155                 sock.connect(format!("127.0.0.1:{}", 4000)).await?;
156                 Arc::new(sock)
157             },
158             payload_type: 111,
159         },
160     );
161     udp_conns.insert(
162         "video".to_owned(),
163         UdpConn {
164             conn: {
165                 let sock = UdpSocket::bind("127.0.0.1:0").await?;
166                 sock.connect(format!("127.0.0.1:{}", 4002)).await?;
167                 Arc::new(sock)
168             },
169             payload_type: 96,
170         },
171     );
172 
173     // Set a handler for when a new remote track starts, this handler will forward data to
174     // our UDP listeners.
175     // In your application this is where you would handle/process audio/video
176     let pc = Arc::downgrade(&peer_connection);
177     peer_connection.on_track(Box::new(move |track, _, _| {
178         // Retrieve udp connection
179         let c = if let Some(c) = udp_conns.get(&track.kind().to_string()) {
180             c.clone()
181         } else {
182             return Box::pin(async {});
183         };
184 
185         // Send a PLI on an interval so that the publisher is pushing a keyframe every rtcpPLIInterval
186         let media_ssrc = track.ssrc();
187         let pc2 = pc.clone();
188         tokio::spawn(async move {
189             let mut result = Result::<usize>::Ok(0);
190             while result.is_ok() {
191                 let timeout = tokio::time::sleep(Duration::from_secs(3));
192                 tokio::pin!(timeout);
193 
194                 tokio::select! {
195                     _ = timeout.as_mut() =>{
196                         if let Some(pc) = pc2.upgrade(){
197                             result = pc.write_rtcp(&[Box::new(PictureLossIndication{
198                                 sender_ssrc: 0,
199                                 media_ssrc,
200                             })]).await.map_err(Into::into);
201                         }else{
202                             break;
203                         }
204                     }
205                 };
206             }
207         });
208 
209         tokio::spawn(async move {
210             let mut b = vec![0u8; 1500];
211             while let Ok((n, _)) = track.read(&mut b).await {
212                 // Unmarshal the packet and update the PayloadType
213                 let mut buf = &b[..n];
214                 let mut rtp_packet = webrtc::rtp::packet::Packet::unmarshal(&mut buf)?;
215                 rtp_packet.header.payload_type = c.payload_type;
216 
217                 // Marshal into original buffer with updated PayloadType
218 
219                 let n = rtp_packet.marshal_to(&mut b)?;
220 
221                 // Write
222                 if let Err(err) = c.conn.send(&b[..n]).await {
223                     // For this particular example, third party applications usually timeout after a short
224                     // amount of time during which the user doesn't have enough time to provide the answer
225                     // to the browser.
226                     // That's why, for this particular example, the user first needs to provide the answer
227                     // to the browser then open the third party application. Therefore we must not kill
228                     // the forward on "connection refused" errors
229                     //if opError, ok := err.(*net.OpError); ok && opError.Err.Error() == "write: connection refused" {
230                     //    continue
231                     //}
232                     //panic(err)
233                     if err.to_string().contains("Connection refused") {
234                         continue;
235                     } else {
236                         println!("conn send err: {err}");
237                         break;
238                     }
239                 }
240             }
241 
242             Result::<()>::Ok(())
243         });
244 
245         Box::pin(async {})
246     }));
247 
248     // Set the handler for ICE connection state
249     // This will notify you when the peer has connected/disconnected
250     peer_connection.on_ice_connection_state_change(Box::new(
251         move |connection_state: RTCIceConnectionState| {
252             println!("Connection State has changed {connection_state}");
253             if connection_state == RTCIceConnectionState::Connected {
254                 println!("Ctrl+C the remote client to stop the demo");
255             }
256             Box::pin(async {})
257         },
258     ));
259 
260     let (done_tx, mut done_rx) = tokio::sync::mpsc::channel::<()>(1);
261 
262     // Set the handler for Peer connection state
263     // This will notify you when the peer has connected/disconnected
264     peer_connection.on_peer_connection_state_change(Box::new(move |s: RTCPeerConnectionState| {
265         println!("Peer Connection State has changed: {s}");
266 
267         if s == RTCPeerConnectionState::Failed {
268             // Wait until PeerConnection has had no network activity for 30 seconds or another failure. It may be reconnected using an ICE Restart.
269             // Use webrtc.PeerConnectionStateDisconnected if you are interested in detecting faster timeout.
270             // Note that the PeerConnection may come back from PeerConnectionStateDisconnected.
271             println!("Peer Connection has gone to failed exiting: Done forwarding");
272             let _ = done_tx.try_send(());
273         }
274 
275         Box::pin(async {})
276     }));
277 
278     // Wait for the offer to be pasted
279     let line = signal::must_read_stdin()?;
280     let desc_data = signal::decode(line.as_str())?;
281     let offer = serde_json::from_str::<RTCSessionDescription>(&desc_data)?;
282 
283     // Set the remote SessionDescription
284     peer_connection.set_remote_description(offer).await?;
285 
286     // Create an answer
287     let answer = peer_connection.create_answer(None).await?;
288 
289     // Create channel that is blocked until ICE Gathering is complete
290     let mut gather_complete = peer_connection.gathering_complete_promise().await;
291 
292     // Sets the LocalDescription, and starts our UDP listeners
293     peer_connection.set_local_description(answer).await?;
294 
295     // Block until ICE Gathering is complete, disabling trickle ICE
296     // we do this because we only can exchange one signaling message
297     // in a production application you should exchange ICE Candidates via OnICECandidate
298     let _ = gather_complete.recv().await;
299 
300     // Output the answer in base64 so we can paste it in browser
301     if let Some(local_desc) = peer_connection.local_description().await {
302         let json_str = serde_json::to_string(&local_desc)?;
303         let b64 = signal::encode(&json_str);
304         println!("{b64}");
305     } else {
306         println!("generate local_description failed!");
307     }
308 
309     println!("Press ctrl-c to stop");
310     tokio::select! {
311         _ = done_rx.recv() => {
312             println!("received done signal!");
313         }
314         _ = tokio::signal::ctrl_c() => {
315             println!();
316         }
317     };
318 
319     peer_connection.close().await?;
320 
321     Ok(())
322 }
323