1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
//! DTLS handshake management
//!
//! This module provides functions for initiating and monitoring DTLS handshakes.
use std::net::SocketAddr;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
use tracing::{debug, error, info, warn};
use crate::api::client::security::dtls::connection;
use crate::api::client::security::srtp::keys;
use crate::api::common::error::SecurityError;
use crate::api::server::security::SocketHandle;
use crate::dtls::DtlsConnection;
use crate::srtp::SrtpContext;
/// Start a DTLS handshake with the remote peer
pub async fn start_handshake(
remote_addr: &Option<SocketAddr>,
connection: &Arc<Mutex<Option<DtlsConnection>>>,
) -> Result<(), SecurityError> {
// Ensure we have a remote address
let remote_addr = remote_addr.ok_or_else(|| {
SecurityError::Configuration("Remote address not set for handshake".to_string())
})?;
// Get the connection
let mut conn_guard = connection.lock().await;
if let Some(conn) = conn_guard.as_mut() {
// Start handshake and send ClientHello
debug!("Starting DTLS handshake with {}", remote_addr);
if let Err(e) = conn.start_handshake(remote_addr).await {
error!("Failed to start DTLS handshake: {}", e);
return Err(SecurityError::Handshake(format!(
"Failed to start DTLS handshake: {}",
e
)));
}
debug!("DTLS handshake started successfully");
Ok(())
} else {
Err(SecurityError::NotInitialized(
"DTLS connection not initialized".to_string(),
))
}
}
/// Wait for a DTLS handshake to complete
pub async fn wait_for_handshake(
connection: &Arc<Mutex<Option<DtlsConnection>>>,
handshake_completed: &Arc<Mutex<bool>>,
srtp_context: &Arc<Mutex<Option<SrtpContext>>>,
) -> Result<(), SecurityError> {
debug!("Waiting for DTLS handshake to complete");
// Get the connection
let mut conn_guard = connection.lock().await;
if let Some(conn) = conn_guard.as_mut() {
// Delegate to the DTLS library's wait_handshake
match conn.wait_handshake().await {
Ok(_) => {
debug!("DTLS handshake completed successfully");
// Set the handshake completed flag
let mut completed = handshake_completed.lock().await;
*completed = true;
// Extract SRTP keys if needed
let srtp_guard = srtp_context.lock().await;
if srtp_guard.is_none() {
// Release the SRTP guard before calling extract_srtp_keys
// to avoid potential deadlock
drop(srtp_guard);
// Extract SRTP keys using the dedicated function
if let Err(e) =
keys::extract_srtp_keys(connection, srtp_context, handshake_completed).await
{
return Err(e);
}
} else {
// We already have SRTP context, just set the completed flag
*completed = true;
}
Ok(())
}
Err(e) => Err(SecurityError::Handshake(format!(
"DTLS handshake failed: {}",
e
))),
}
} else {
Err(SecurityError::NotInitialized(
"DTLS connection not initialized".to_string(),
))
}
}
/// Start a handshake monitor task
pub async fn start_handshake_monitor(
handshake_monitor_running: &Arc<AtomicBool>,
remote_addr: &Arc<Mutex<Option<SocketAddr>>>,
socket: &Arc<Mutex<Option<SocketHandle>>>,
connection: &Arc<Mutex<Option<DtlsConnection>>>,
handshake_completed: &Arc<Mutex<bool>>,
) -> Result<(), SecurityError> {
// If already running, don't start another one
if handshake_monitor_running.load(Ordering::SeqCst) {
return Ok(());
}
debug!("Starting handshake monitoring task");
// Set running flag
handshake_monitor_running.store(true, Ordering::SeqCst);
// Get what we need for the task - remote address and socket
let remote_addr_guard = remote_addr.lock().await;
let remote_addr = match remote_addr_guard.as_ref() {
Some(addr) => *addr,
None => {
handshake_monitor_running.store(false, Ordering::SeqCst);
return Err(SecurityError::Configuration(
"Remote address not set for handshake monitor".to_string(),
));
}
};
drop(remote_addr_guard);
let socket_guard = socket.lock().await;
let socket = match socket_guard.as_ref() {
Some(s) => s.socket.clone(),
None => {
handshake_monitor_running.store(false, Ordering::SeqCst);
return Err(SecurityError::Configuration(
"Socket not set for handshake monitor".to_string(),
));
}
};
drop(socket_guard);
// Clone all the required Arc references for the task
let connection_ref = connection.clone();
let handshake_completed_ref = handshake_completed.clone();
let handshake_monitor_running_ref = handshake_monitor_running.clone();
// Spawn the monitor task
tokio::spawn(async move {
debug!("Handshake monitor task started");
let mut attempt = 1;
let max_attempts = 3;
// Try the handshake process multiple times
while attempt <= max_attempts {
debug!("Handshake attempt #{}", attempt);
// First check if handshake is already complete
if *handshake_completed_ref.lock().await {
debug!("Handshake already completed, stopping monitor");
handshake_monitor_running_ref.store(false, Ordering::SeqCst);
return;
}
// Wait to see if the handshake completes in a reasonable time
let mut completed = false;
for _ in 0..30 {
// Check status every 100ms for up to 3 seconds
// Check handshake status
let conn_guard = connection_ref.lock().await;
if let Some(conn) = conn_guard.as_ref() {
// Check connection state first
if conn.state() == crate::dtls::connection::ConnectionState::Connected {
debug!("Connection state is Connected, handshake likely complete");
completed = true;
break;
}
// Check handshake step (it may need assistance)
if let Some(step) = conn.handshake_step() {
match step {
crate::dtls::handshake::HandshakeStep::ReceivedHelloVerifyRequest => {
// This requires an explicit restart of handshake with cookie
debug!("Detected ReceivedHelloVerifyRequest step - need to restart handshake");
drop(conn_guard);
// Get fresh connection guard for mutation
let mut conn_guard = connection_ref.lock().await;
if let Some(conn) = conn_guard.as_mut() {
// HelloVerifyRequest received, restart handshake with cookie
debug!("Restarting handshake with cookie from monitor");
if let Err(e) = conn.start_handshake(remote_addr).await {
warn!("Failed to restart handshake with cookie: {}", e);
} else {
info!("Successfully restarted handshake with cookie");
}
}
break; // Exit loop to allow handshake to progress
}
crate::dtls::handshake::HandshakeStep::Complete => {
debug!("Handshake step is Complete");
completed = true;
break;
}
crate::dtls::handshake::HandshakeStep::Failed => {
debug!("Handshake step is Failed, will restart");
break;
}
_ => {
// Continue waiting
debug!("Waiting for handshake progress, current step: {:?}", step);
}
}
}
}
// Check for completion flag
if *handshake_completed_ref.lock().await {
debug!("Handshake completed flag is set");
completed = true;
break;
}
// Wait a bit before checking again
tokio::time::sleep(Duration::from_millis(100)).await;
}
// If completed, we're done
if completed || *handshake_completed_ref.lock().await {
debug!("Handshake completed successfully, exiting monitor");
handshake_monitor_running_ref.store(false, Ordering::SeqCst);
return;
}
// If we get here, we need to try again with a new connection
warn!(
"Creating new DTLS connection for retry (attempt #{})",
attempt + 1
);
// Create a new connection and transport
match connection::create_connection(&socket, remote_addr).await {
Ok(new_conn) => {
// Replace the old connection
let mut conn_guard = connection_ref.lock().await;
*conn_guard = Some(new_conn);
}
Err(e) => {
error!("Failed to create new connection: {}", e);
}
}
// Increment attempt counter
attempt += 1;
// Wait a bit before next attempt
tokio::time::sleep(Duration::from_millis(500)).await;
}
warn!("Handshake failed after {} attempts", max_attempts);
handshake_monitor_running_ref.store(false, Ordering::SeqCst);
});
Ok(())
}
/// Check if a DTLS handshake is complete
pub fn is_handshake_complete(handshake_completed: &Arc<Mutex<bool>>) -> bool {
// Since this is a sync function, we need to block on the async operation
// We could use a blocking_lock() here, but that could lead to deadlocks
// Instead, we'll just default to false if we can't get the lock immediately
match handshake_completed.try_lock() {
Ok(guard) => *guard,
Err(_) => false,
}
}