1 use anyhow::Result;
2 use clap::{AppSettings, Arg, Command};
3 use std::io::Write;
4 use std::sync::atomic::{AtomicI32, Ordering};
5 use std::sync::Arc;
6 use tokio::sync::Mutex;
7 use tokio::time::Duration;
8 use webrtc::api::interceptor_registry::register_default_interceptors;
9 use webrtc::api::media_engine::MediaEngine;
10 use webrtc::api::APIBuilder;
11 use webrtc::data_channel::data_channel_message::DataChannelMessage;
12 use webrtc::data_channel::RTCDataChannel;
13 use webrtc::ice_transport::ice_server::RTCIceServer;
14 use webrtc::interceptor::registry::Registry;
15 use webrtc::peer_connection::configuration::RTCConfiguration;
16 use webrtc::peer_connection::math_rand_alpha;
17 use webrtc::peer_connection::peer_connection_state::RTCPeerConnectionState;
18 use webrtc::peer_connection::sdp::session_description::RTCSessionDescription;
19 
20 #[tokio::main]
main() -> Result<()>21 async fn main() -> Result<()> {
22     let mut app = Command::new("data-channels-close")
23         .version("0.1.0")
24         .author("Rain Liu <[email protected]>")
25         .about("An example of Data-Channels-Close.")
26         .setting(AppSettings::DeriveDisplayOrder)
27         .subcommand_negates_reqs(true)
28         .arg(
29             Arg::new("FULLHELP")
30                 .help("Prints more detailed help information")
31                 .long("fullhelp"),
32         )
33         .arg(
34             Arg::new("debug")
35                 .long("debug")
36                 .short('d')
37                 .help("Prints debug log information"),
38         )
39         .arg(
40             Arg::new("close-after")
41                 .takes_value(true)
42                 .default_value("5")
43                 .long("close-after")
44                 .help("Close data channel after sending X times."),
45         );
46 
47     let matches = app.clone().get_matches();
48 
49     if matches.is_present("FULLHELP") {
50         app.print_long_help().unwrap();
51         std::process::exit(0);
52     }
53 
54     let close_after = Arc::new(AtomicI32::new(
55         matches
56             .value_of("close-after")
57             .unwrap()
58             .to_owned()
59             .parse::<i32>()?,
60     ));
61     let debug = matches.is_present("debug");
62     if debug {
63         env_logger::Builder::new()
64             .format(|buf, record| {
65                 writeln!(
66                     buf,
67                     "{}:{} [{}] {} - {}",
68                     record.file().unwrap_or("unknown"),
69                     record.line().unwrap_or(0),
70                     record.level(),
71                     chrono::Local::now().format("%H:%M:%S.%6f"),
72                     record.args()
73                 )
74             })
75             .filter(None, log::LevelFilter::Trace)
76             .init();
77     }
78 
79     // Everything below is the WebRTC-rs API! Thanks for using it ❤️.
80 
81     // Create a MediaEngine object to configure the supported codec
82     let mut m = MediaEngine::default();
83 
84     // Register default codecs
85     m.register_default_codecs()?;
86 
87     // Create a InterceptorRegistry. This is the user configurable RTP/RTCP Pipeline.
88     // This provides NACKs, RTCP Reports and other features. If you use `webrtc.NewPeerConnection`
89     // this is enabled by default. If you are manually managing You MUST create a InterceptorRegistry
90     // for each PeerConnection.
91     let mut registry = Registry::new();
92 
93     // Use the default set of Interceptors
94     registry = register_default_interceptors(registry, &mut m)?;
95 
96     // Create the API object with the MediaEngine
97     let api = APIBuilder::new()
98         .with_media_engine(m)
99         .with_interceptor_registry(registry)
100         .build();
101 
102     // Prepare the configuration
103     let config = RTCConfiguration {
104         ice_servers: vec![RTCIceServer {
105             urls: vec!["stun:stun.l.google.com:19302".to_owned()],
106             ..Default::default()
107         }],
108         ..Default::default()
109     };
110 
111     // Create a new RTCPeerConnection
112     let peer_connection = Arc::new(api.new_peer_connection(config).await?);
113 
114     let (done_tx, mut done_rx) = tokio::sync::mpsc::channel::<()>(1);
115 
116     // Set the handler for Peer connection state
117     // This will notify you when the peer has connected/disconnected
118     peer_connection.on_peer_connection_state_change(Box::new(move |s: RTCPeerConnectionState| {
119         println!("Peer Connection State has changed: {s}");
120 
121         if s == RTCPeerConnectionState::Failed {
122             // Wait until PeerConnection has had no network activity for 30 seconds or another failure. It may be reconnected using an ICE Restart.
123             // Use webrtc.PeerConnectionStateDisconnected if you are interested in detecting faster timeout.
124             // Note that the PeerConnection may come back from PeerConnectionStateDisconnected.
125             println!("Peer Connection has gone to failed exiting");
126             let _ = done_tx.try_send(());
127         }
128 
129         Box::pin(async {})
130     }));
131 
132     // Register data channel creation handling
133     peer_connection
134         .on_data_channel(Box::new(move |d: Arc<RTCDataChannel>| {
135             let d_label = d.label().to_owned();
136             let d_id = d.id();
137             println!("New DataChannel {d_label} {d_id}");
138 
139             let close_after2 = Arc::clone(&close_after);
140 
141             // Register channel opening handling
142             Box::pin(async move {
143                 let d2 = Arc::clone(&d);
144                 let d_label2 = d_label.clone();
145                 let d_id2 = d_id;
146                 d.on_open(Box::new(move || {
147                     println!("Data channel '{d_label2}'-'{d_id2}' open. Random messages will now be sent to any connected DataChannels every 5 seconds");
148                     let (done_tx, mut done_rx) = tokio::sync::mpsc::channel::<()>(1);
149                     let done_tx = Arc::new(Mutex::new(Some(done_tx)));
150                     Box::pin(async move {
151                         d2.on_close(Box::new(move || {
152                             println!("Data channel '{d_label2}'-'{d_id2}' closed.");
153                             let done_tx2 = Arc::clone(&done_tx);
154                             Box::pin(async move{
155                                 let mut done = done_tx2.lock().await;
156                                 done.take();
157                             })
158                         }));
159 
160                         let mut result = Result::<usize>::Ok(0);
161                         while result.is_ok() {
162                             let timeout = tokio::time::sleep(Duration::from_secs(5));
163                             tokio::pin!(timeout);
164 
165                             tokio::select! {
166                                 _ = done_rx.recv() => {
167                                     break;
168                                 }
169                                 _ = timeout.as_mut() =>{
170                                     let message = math_rand_alpha(15);
171                                     println!("Sending '{message}'");
172                                     result = d2.send_text(message).await.map_err(Into::into);
173 
174                                     let cnt = close_after2.fetch_sub(1, Ordering::SeqCst);
175                                     if cnt <= 0 {
176                                         println!("Sent times out. Closing data channel '{}'-'{}'.", d2.label(), d2.id());
177                                         let _ = d2.close().await;
178                                         break;
179                                     }
180                                 }
181                             };
182                         }
183                     })
184                 }));
185 
186                 // Register text message handling
187                 d.on_message(Box::new(move |msg: DataChannelMessage| {
188                     let msg_str = String::from_utf8(msg.data.to_vec()).unwrap();
189                     println!("Message from DataChannel '{d_label}': '{msg_str}'");
190                     Box::pin(async {})
191                 }));
192             })
193         }));
194 
195     // Wait for the offer to be pasted
196     let line = signal::must_read_stdin()?;
197     let desc_data = signal::decode(line.as_str())?;
198     let offer = serde_json::from_str::<RTCSessionDescription>(&desc_data)?;
199 
200     // Set the remote SessionDescription
201     peer_connection.set_remote_description(offer).await?;
202 
203     // Create an answer
204     let answer = peer_connection.create_answer(None).await?;
205 
206     // Create channel that is blocked until ICE Gathering is complete
207     let mut gather_complete = peer_connection.gathering_complete_promise().await;
208 
209     // Sets the LocalDescription, and starts our UDP listeners
210     peer_connection.set_local_description(answer).await?;
211 
212     // Block until ICE Gathering is complete, disabling trickle ICE
213     // we do this because we only can exchange one signaling message
214     // in a production application you should exchange ICE Candidates via OnICECandidate
215     let _ = gather_complete.recv().await;
216 
217     // Output the answer in base64 so we can paste it in browser
218     if let Some(local_desc) = peer_connection.local_description().await {
219         let json_str = serde_json::to_string(&local_desc)?;
220         let b64 = signal::encode(&json_str);
221         println!("{b64}");
222     } else {
223         println!("generate local_description failed!");
224     }
225 
226     println!("Press ctrl-c to stop");
227     tokio::select! {
228         _ = done_rx.recv() => {
229             println!("received done signal!");
230         }
231         _ = tokio::signal::ctrl_c() => {
232             println!();
233         }
234     };
235 
236     peer_connection.close().await?;
237 
238     Ok(())
239 }
240