1use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
8use serde::{Deserialize, Serialize};
9use serde_json::Value;
10use std::collections::HashMap;
11use std::sync::{Arc, Mutex};
12use tungstenite::{connect, Message};
13use url::Url;
14
15use crate::error::{Error, Result};
16use crate::model::HttpRequest;
17
18pub mod phase {
24 pub const REQUEST: &str = "REQUEST";
25 pub const RESPONSE: &str = "RESPONSE";
26 pub const RESPONSE_STREAM: &str = "RESPONSE_STREAM";
27 pub const INBOUND_STREAM: &str = "INBOUND_STREAM";
28}
29
30#[derive(Debug, Clone, Serialize, Deserialize)]
36#[serde(rename_all = "camelCase")]
37pub struct BreakpointMatcherRegistration {
38 pub http_request: HttpRequest,
39 pub phases: Vec<String>,
40 #[serde(skip_serializing_if = "Option::is_none")]
41 pub client_id: Option<String>,
42}
43
44#[derive(Debug, Clone, Serialize, Deserialize)]
46pub struct BreakpointMatcherResponse {
47 pub id: String,
48 pub phases: Vec<String>,
49}
50
51#[derive(Debug, Clone, Serialize, Deserialize)]
53#[serde(rename_all = "camelCase")]
54pub struct BreakpointMatcherEntry {
55 pub id: String,
56 pub http_request: Value,
57 pub phases: Vec<String>,
58 #[serde(skip_serializing_if = "Option::is_none")]
59 pub client_id: Option<String>,
60}
61
62#[derive(Debug, Clone, Serialize, Deserialize)]
64pub struct BreakpointMatcherList {
65 pub matchers: Vec<BreakpointMatcherEntry>,
66}
67
68#[derive(Debug, Clone, Serialize, Deserialize)]
70#[serde(rename_all = "camelCase")]
71pub struct PausedStreamFrame {
72 pub correlation_id: String,
73 #[serde(default)]
74 pub stream_id: String,
75 #[serde(default)]
76 pub sequence_number: i64,
77 #[serde(default)]
78 pub direction: String,
79 #[serde(default)]
80 pub phase: String,
81 #[serde(default)]
82 pub body: String,
83 #[serde(skip_serializing_if = "Option::is_none")]
84 pub request_method: Option<String>,
85 #[serde(skip_serializing_if = "Option::is_none")]
86 pub request_path: Option<String>,
87 #[serde(skip_serializing_if = "Option::is_none")]
88 pub breakpoint_id: Option<String>,
89}
90
91impl PausedStreamFrame {
92 pub fn body_bytes(&self) -> std::result::Result<Vec<u8>, base64::DecodeError> {
94 BASE64.decode(&self.body)
95 }
96}
97
98#[derive(Debug, Clone, Serialize, Deserialize)]
100#[serde(rename_all = "camelCase")]
101pub struct StreamFrameDecision {
102 pub correlation_id: String,
103 pub action: String,
104 #[serde(skip_serializing_if = "Option::is_none")]
105 pub body: Option<String>,
106}
107
108impl StreamFrameDecision {
109 pub fn continue_frame(correlation_id: impl Into<String>) -> Self {
111 Self {
112 correlation_id: correlation_id.into(),
113 action: "CONTINUE".to_string(),
114 body: None,
115 }
116 }
117
118 pub fn modify(correlation_id: impl Into<String>, body: &[u8]) -> Self {
120 Self {
121 correlation_id: correlation_id.into(),
122 action: "MODIFY".to_string(),
123 body: Some(BASE64.encode(body)),
124 }
125 }
126
127 pub fn drop_frame(correlation_id: impl Into<String>) -> Self {
129 Self {
130 correlation_id: correlation_id.into(),
131 action: "DROP".to_string(),
132 body: None,
133 }
134 }
135
136 pub fn inject(correlation_id: impl Into<String>, extra_body: &[u8]) -> Self {
138 Self {
139 correlation_id: correlation_id.into(),
140 action: "INJECT".to_string(),
141 body: Some(BASE64.encode(extra_body)),
142 }
143 }
144
145 pub fn close(correlation_id: impl Into<String>) -> Self {
147 Self {
148 correlation_id: correlation_id.into(),
149 action: "CLOSE".to_string(),
150 body: None,
151 }
152 }
153}
154
155#[derive(Debug, Clone, Serialize, Deserialize)]
160pub struct WsEnvelope {
161 #[serde(rename = "type")]
162 pub msg_type: String,
163 pub value: String,
164}
165
166#[derive(Debug, Clone, Deserialize)]
167#[serde(rename_all = "camelCase")]
168struct WsClientIdDTO {
169 client_id: String,
170}
171
172pub type BreakpointRequestHandler = Box<dyn Fn(Value) -> Option<Value> + Send + Sync>;
183
184pub type BreakpointResponseHandler = Box<dyn Fn(Value, Value) -> Option<Value> + Send + Sync>;
188
189pub type BreakpointStreamFrameHandler =
193 Box<dyn Fn(&PausedStreamFrame) -> Option<StreamFrameDecision> + Send + Sync>;
194
195pub type ObjectResponseHandler = Box<dyn Fn(Value) -> Value + Send>;
207
208pub(crate) struct BreakpointWebSocketClient {
214 pub(crate) client_id: String,
215 socket: Arc<Mutex<tungstenite::WebSocket<tungstenite::stream::MaybeTlsStream<std::net::TcpStream>>>>,
216 request_handlers: Arc<Mutex<HashMap<String, BreakpointRequestHandler>>>,
217 response_handlers: Arc<Mutex<HashMap<String, BreakpointResponseHandler>>>,
218 stream_frame_handlers: Arc<Mutex<HashMap<String, BreakpointStreamFrameHandler>>>,
219 object_response_handler: Arc<Mutex<Option<ObjectResponseHandler>>>,
222 dead: Arc<std::sync::atomic::AtomicBool>,
223 _read_thread: Option<std::thread::JoinHandle<()>>,
224}
225
226impl BreakpointWebSocketClient {
227 pub(crate) fn is_dead(&self) -> bool {
229 self.dead.load(std::sync::atomic::Ordering::Relaxed)
230 }
231
232 pub(crate) fn connect(base_url: &str) -> Result<Self> {
234 let parsed = Url::parse(base_url)
235 .map_err(|e| Error::InvalidRequest(format!("bad base URL: {e}")))?;
236
237 let ws_scheme = if parsed.scheme() == "https" { "wss" } else { "ws" };
238 let host = parsed.host_str().unwrap_or("localhost");
239 let port = parsed.port().unwrap_or(if parsed.scheme() == "https" { 443 } else { 80 });
240 let path = parsed.path().trim_end_matches('/');
241 let ws_url = format!("{ws_scheme}://{host}:{port}{path}/_mockserver_callback_websocket");
242
243 let (mut socket, _response) = connect(&ws_url)
249 .map_err(|e| Error::InvalidRequest(format!("WebSocket connect failed: {e}")))?;
250
251 let reg_msg = socket
253 .read()
254 .map_err(|e| Error::InvalidRequest(format!("WebSocket read registration: {e}")))?;
255
256 let reg_text = match reg_msg {
257 Message::Text(t) => t,
258 _ => return Err(Error::InvalidRequest("Expected text registration message".into())),
259 };
260
261 let envelope: WsEnvelope = serde_json::from_str(®_text)?;
262 let client_id = if envelope.msg_type
263 == "org.mockserver.serialization.model.WebSocketClientIdDTO"
264 {
265 let dto: WsClientIdDTO = serde_json::from_str(&envelope.value)?;
266 dto.client_id
267 } else {
268 return Err(Error::InvalidRequest(format!(
269 "Unexpected registration type: {}",
270 envelope.msg_type
271 )));
272 };
273
274 let request_handlers: Arc<Mutex<HashMap<String, BreakpointRequestHandler>>> =
275 Arc::new(Mutex::new(HashMap::new()));
276 let response_handlers: Arc<Mutex<HashMap<String, BreakpointResponseHandler>>> =
277 Arc::new(Mutex::new(HashMap::new()));
278 let stream_frame_handlers: Arc<Mutex<HashMap<String, BreakpointStreamFrameHandler>>> =
279 Arc::new(Mutex::new(HashMap::new()));
280 let object_response_handler: Arc<Mutex<Option<ObjectResponseHandler>>> =
281 Arc::new(Mutex::new(None));
282
283 let socket = Arc::new(Mutex::new(socket));
284 let dead = Arc::new(std::sync::atomic::AtomicBool::new(false));
285
286 let rh = Arc::clone(&request_handlers);
288 let resh = Arc::clone(&response_handlers);
289 let sfh = Arc::clone(&stream_frame_handlers);
290 let orh = Arc::clone(&object_response_handler);
291 let sock = Arc::clone(&socket);
292 let dead_flag = Arc::clone(&dead);
293
294 let read_thread = std::thread::spawn(move || {
295 read_loop(sock, rh, resh, sfh, orh, dead_flag);
296 });
297
298 Ok(Self {
299 client_id,
300 socket,
301 request_handlers,
302 response_handlers,
303 stream_frame_handlers,
304 object_response_handler,
305 dead,
306 _read_thread: Some(read_thread),
307 })
308 }
309
310 pub(crate) fn set_request_handler(&self, breakpoint_id: &str, handler: BreakpointRequestHandler) {
311 self.request_handlers
312 .lock()
313 .unwrap()
314 .insert(breakpoint_id.to_string(), handler);
315 }
316
317 pub(crate) fn set_response_handler(&self, breakpoint_id: &str, handler: BreakpointResponseHandler) {
318 self.response_handlers
319 .lock()
320 .unwrap()
321 .insert(breakpoint_id.to_string(), handler);
322 }
323
324 pub(crate) fn set_stream_frame_handler(
325 &self,
326 breakpoint_id: &str,
327 handler: BreakpointStreamFrameHandler,
328 ) {
329 self.stream_frame_handlers
330 .lock()
331 .unwrap()
332 .insert(breakpoint_id.to_string(), handler);
333 }
334
335 pub(crate) fn set_object_response_handler(&self, handler: ObjectResponseHandler) {
337 *self.object_response_handler.lock().unwrap() = Some(handler);
338 }
339
340 pub(crate) fn remove_handlers(&self, breakpoint_id: &str) {
341 self.request_handlers.lock().unwrap().remove(breakpoint_id);
342 self.response_handlers.lock().unwrap().remove(breakpoint_id);
343 self.stream_frame_handlers
344 .lock()
345 .unwrap()
346 .remove(breakpoint_id);
347 }
348
349 pub(crate) fn clear_handlers(&self) {
350 self.request_handlers.lock().unwrap().clear();
351 self.response_handlers.lock().unwrap().clear();
352 self.stream_frame_handlers.lock().unwrap().clear();
353 }
354
355 pub(crate) fn close(&self) {
356 if let Ok(mut sock) = self.socket.lock() {
357 let _ = sock.close(None);
358 }
359 }
360}
361
362struct ReadLoopDeadGuard(Arc<std::sync::atomic::AtomicBool>);
368impl Drop for ReadLoopDeadGuard {
369 fn drop(&mut self) {
370 self.0.store(true, std::sync::atomic::Ordering::Relaxed);
371 }
372}
373
374#[allow(clippy::too_many_arguments)]
375fn read_loop(
376 socket: Arc<Mutex<tungstenite::WebSocket<tungstenite::stream::MaybeTlsStream<std::net::TcpStream>>>>,
377 request_handlers: Arc<Mutex<HashMap<String, BreakpointRequestHandler>>>,
378 response_handlers: Arc<Mutex<HashMap<String, BreakpointResponseHandler>>>,
379 stream_frame_handlers: Arc<Mutex<HashMap<String, BreakpointStreamFrameHandler>>>,
380 object_response_handler: Arc<Mutex<Option<ObjectResponseHandler>>>,
381 dead: Arc<std::sync::atomic::AtomicBool>,
382) {
383 let _guard = ReadLoopDeadGuard(dead);
384 loop {
385 let msg = {
386 let mut sock = match socket.lock() {
387 Ok(s) => s,
388 Err(e) => {
389 eprintln!("mockserver: breakpoint ws read loop terminated: mutex poisoned: {e}");
390 return;
391 }
392 };
393 match sock.read() {
394 Ok(m) => m,
395 Err(e) => {
396 eprintln!("mockserver: breakpoint ws read loop terminated: {e}");
397 return;
398 }
399 }
400 };
401
402 let text = match msg {
403 Message::Text(t) => t,
404 Message::Close(_) => {
405 eprintln!("mockserver: breakpoint ws read loop terminated: received close frame");
406 return;
407 }
408 _ => continue,
409 };
410
411 let envelope: WsEnvelope = match serde_json::from_str(&text) {
412 Ok(e) => e,
413 Err(_) => continue,
414 };
415
416 match envelope.msg_type.as_str() {
417 "org.mockserver.model.HttpRequest" => {
418 handle_request(
419 &envelope.value,
420 &request_handlers,
421 &object_response_handler,
422 &socket,
423 );
424 }
425 "org.mockserver.model.HttpRequestAndHttpResponse" => {
426 handle_response(
427 &envelope.value,
428 &response_handlers,
429 &socket,
430 );
431 }
432 "org.mockserver.serialization.model.PausedStreamFrameDTO" => {
433 handle_stream_frame(
434 &envelope.value,
435 &stream_frame_handlers,
436 &socket,
437 );
438 }
439 "org.mockserver.serialization.model.WebSocketClientIdDTO" => {
440 }
442 _ => {}
443 }
444 }
445}
446
447fn handle_request(
448 value_json: &str,
449 handlers: &Arc<Mutex<HashMap<String, BreakpointRequestHandler>>>,
450 object_response_handler: &Arc<Mutex<Option<ObjectResponseHandler>>>,
451 socket: &Arc<Mutex<tungstenite::WebSocket<tungstenite::stream::MaybeTlsStream<std::net::TcpStream>>>>,
452) {
453 if let Some((type_name, result)) = route_object_callback(value_json, object_response_handler) {
457 send_envelope(socket, &type_name, &result);
458 return;
459 }
460 if let Some((type_name, result)) = route_request(value_json, handlers) {
461 send_envelope(socket, &type_name, &result);
462 }
463}
464
465fn handle_response(
466 value_json: &str,
467 handlers: &Arc<Mutex<HashMap<String, BreakpointResponseHandler>>>,
468 socket: &Arc<Mutex<tungstenite::WebSocket<tungstenite::stream::MaybeTlsStream<std::net::TcpStream>>>>,
469) {
470 if let Some((type_name, result)) = route_response(value_json, handlers) {
471 send_envelope(socket, &type_name, &result);
472 }
473}
474
475fn handle_stream_frame(
476 value_json: &str,
477 handlers: &Arc<Mutex<HashMap<String, BreakpointStreamFrameHandler>>>,
478 socket: &Arc<Mutex<tungstenite::WebSocket<tungstenite::stream::MaybeTlsStream<std::net::TcpStream>>>>,
479) {
480 if let Some((type_name, decision)) = route_stream_frame(value_json, handlers) {
481 send_envelope(socket, &type_name, &decision);
482 }
483}
484
485pub fn route_request(
492 value_json: &str,
493 handlers: &Arc<Mutex<HashMap<String, BreakpointRequestHandler>>>,
494) -> Option<(String, Value)> {
495 let request: Value = serde_json::from_str(value_json).ok()?;
496
497 let correlation_id = extract_header(&request, "WebSocketCorrelationId");
498 let breakpoint_id = extract_header(&request, "X-MockServer-BreakpointId");
499
500 let mut result: Option<Value> = None;
501 if let Some(bp_id) = &breakpoint_id {
502 if let Ok(h) = handlers.lock() {
503 if let Some(handler) = h.get(bp_id) {
504 result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
505 handler(request.clone())
506 }))
507 .unwrap_or(None);
508 }
509 }
510 }
511
512 let mut result = result.unwrap_or_else(|| request.clone());
514
515 let type_name = if result.get("statusCode").is_some() {
517 "org.mockserver.model.HttpResponse"
518 } else {
519 "org.mockserver.model.HttpRequest"
520 };
521
522 if let Some(corr) = &correlation_id {
523 set_header(&mut result, "WebSocketCorrelationId", corr);
524 }
525
526 Some((type_name.to_string(), result))
527}
528
529pub fn route_object_callback(
542 value_json: &str,
543 handler: &Arc<Mutex<Option<ObjectResponseHandler>>>,
544) -> Option<(String, Value)> {
545 let request: Value = serde_json::from_str(value_json).ok()?;
546
547 if extract_header(&request, "X-MockServer-BreakpointId").is_some() {
549 return None;
550 }
551
552 let correlation_id = extract_header(&request, "WebSocketCorrelationId");
553
554 let guard = handler.lock().ok()?;
556 let handler = guard.as_ref()?;
557
558 let mut result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
559 handler(request.clone())
560 }))
561 .unwrap_or_else(|_| serde_json::json!({}));
562
563 if let Some(corr) = &correlation_id {
564 set_header(&mut result, "WebSocketCorrelationId", corr);
565 }
566
567 Some(("org.mockserver.model.HttpResponse".to_string(), result))
568}
569
570pub fn route_response(
573 value_json: &str,
574 handlers: &Arc<Mutex<HashMap<String, BreakpointResponseHandler>>>,
575) -> Option<(String, Value)> {
576 let req_and_resp: Value = serde_json::from_str(value_json).ok()?;
577
578 let http_request = req_and_resp.get("httpRequest").cloned().unwrap_or(Value::Null);
579 let http_response = req_and_resp.get("httpResponse").cloned().unwrap_or(Value::Null);
580
581 let correlation_id = extract_header(&http_request, "WebSocketCorrelationId");
582 let breakpoint_id = extract_header(&http_request, "X-MockServer-BreakpointId");
583
584 let mut result: Option<Value> = None;
585 if let Some(bp_id) = &breakpoint_id {
586 if let Ok(h) = handlers.lock() {
587 if let Some(handler) = h.get(bp_id) {
588 result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
589 handler(http_request.clone(), http_response.clone())
590 }))
591 .unwrap_or(None);
592 }
593 }
594 }
595
596 let mut result = result.unwrap_or(http_response);
597
598 if let Some(corr) = &correlation_id {
599 set_header(&mut result, "WebSocketCorrelationId", corr);
600 }
601
602 Some(("org.mockserver.model.HttpResponse".to_string(), result))
603}
604
605pub fn route_stream_frame(
608 value_json: &str,
609 handlers: &Arc<Mutex<HashMap<String, BreakpointStreamFrameHandler>>>,
610) -> Option<(String, StreamFrameDecision)> {
611 let frame: PausedStreamFrame = serde_json::from_str(value_json).ok()?;
612
613 let mut decision: Option<StreamFrameDecision> = None;
614 if let Some(bp_id) = &frame.breakpoint_id {
615 if let Ok(h) = handlers.lock() {
616 if let Some(handler) = h.get(bp_id) {
617 decision = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
618 handler(&frame)
619 }))
620 .unwrap_or(None);
621 }
622 }
623 }
624
625 let mut decision = decision.unwrap_or_else(|| StreamFrameDecision::continue_frame(&frame.correlation_id));
626 decision.correlation_id = frame.correlation_id;
627
628 Some((
629 "org.mockserver.serialization.model.StreamFrameDecisionDTO".to_string(),
630 decision,
631 ))
632}
633
634fn send_envelope<T: Serialize>(
635 socket: &Arc<Mutex<tungstenite::WebSocket<tungstenite::stream::MaybeTlsStream<std::net::TcpStream>>>>,
636 type_name: &str,
637 value: &T,
638) {
639 let value_json = match serde_json::to_string(value) {
640 Ok(j) => j,
641 Err(_) => return,
642 };
643 let envelope = WsEnvelope {
644 msg_type: type_name.to_string(),
645 value: value_json,
646 };
647 let env_json = match serde_json::to_string(&envelope) {
648 Ok(j) => j,
649 Err(_) => return,
650 };
651 if let Ok(mut sock) = socket.lock() {
652 let _ = sock.send(Message::Text(env_json));
653 }
654}
655
656pub fn extract_header(obj: &Value, name: &str) -> Option<String> {
662 let headers = obj.get("headers")?.as_object()?;
663 let name_lower = name.to_lowercase();
664 for (k, v) in headers {
665 if k.to_lowercase() == name_lower {
666 if let Some(arr) = v.as_array() {
667 return arr.first()?.as_str().map(|s| s.to_string());
668 }
669 return v.as_str().map(|s| s.to_string());
670 }
671 }
672 None
673}
674
675pub fn set_header(obj: &mut Value, name: &str, value: &str) {
677 if let Some(obj_map) = obj.as_object_mut() {
678 let headers = obj_map
679 .entry("headers")
680 .or_insert_with(|| Value::Object(serde_json::Map::new()));
681 if let Some(h) = headers.as_object_mut() {
682 h.insert(
683 name.to_string(),
684 serde_json::json!([value]),
685 );
686 }
687 }
688}