use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use tungstenite::{connect, Message};
use url::Url;
use crate::error::{Error, Result};
use crate::model::HttpRequest;
pub mod phase {
pub const REQUEST: &str = "REQUEST";
pub const RESPONSE: &str = "RESPONSE";
pub const RESPONSE_STREAM: &str = "RESPONSE_STREAM";
pub const INBOUND_STREAM: &str = "INBOUND_STREAM";
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct BreakpointMatcherRegistration {
pub http_request: HttpRequest,
pub phases: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub client_id: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BreakpointMatcherResponse {
pub id: String,
pub phases: Vec<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct BreakpointMatcherEntry {
pub id: String,
pub http_request: Value,
pub phases: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub client_id: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BreakpointMatcherList {
pub matchers: Vec<BreakpointMatcherEntry>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PausedStreamFrame {
pub correlation_id: String,
#[serde(default)]
pub stream_id: String,
#[serde(default)]
pub sequence_number: i64,
#[serde(default)]
pub direction: String,
#[serde(default)]
pub phase: String,
#[serde(default)]
pub body: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub request_method: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub request_path: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub breakpoint_id: Option<String>,
}
impl PausedStreamFrame {
pub fn body_bytes(&self) -> std::result::Result<Vec<u8>, base64::DecodeError> {
BASE64.decode(&self.body)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct StreamFrameDecision {
pub correlation_id: String,
pub action: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub body: Option<String>,
}
impl StreamFrameDecision {
pub fn continue_frame(correlation_id: impl Into<String>) -> Self {
Self {
correlation_id: correlation_id.into(),
action: "CONTINUE".to_string(),
body: None,
}
}
pub fn modify(correlation_id: impl Into<String>, body: &[u8]) -> Self {
Self {
correlation_id: correlation_id.into(),
action: "MODIFY".to_string(),
body: Some(BASE64.encode(body)),
}
}
pub fn drop_frame(correlation_id: impl Into<String>) -> Self {
Self {
correlation_id: correlation_id.into(),
action: "DROP".to_string(),
body: None,
}
}
pub fn inject(correlation_id: impl Into<String>, extra_body: &[u8]) -> Self {
Self {
correlation_id: correlation_id.into(),
action: "INJECT".to_string(),
body: Some(BASE64.encode(extra_body)),
}
}
pub fn close(correlation_id: impl Into<String>) -> Self {
Self {
correlation_id: correlation_id.into(),
action: "CLOSE".to_string(),
body: None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WsEnvelope {
#[serde(rename = "type")]
pub msg_type: String,
pub value: String,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct WsClientIdDTO {
client_id: String,
}
pub type BreakpointRequestHandler = Box<dyn Fn(Value) -> Option<Value> + Send + Sync>;
pub type BreakpointResponseHandler = Box<dyn Fn(Value, Value) -> Option<Value> + Send + Sync>;
pub type BreakpointStreamFrameHandler =
Box<dyn Fn(&PausedStreamFrame) -> Option<StreamFrameDecision> + Send + Sync>;
pub type ObjectResponseHandler = Box<dyn Fn(Value) -> Value + Send>;
pub(crate) struct BreakpointWebSocketClient {
pub(crate) client_id: String,
socket: Arc<Mutex<tungstenite::WebSocket<tungstenite::stream::MaybeTlsStream<std::net::TcpStream>>>>,
request_handlers: Arc<Mutex<HashMap<String, BreakpointRequestHandler>>>,
response_handlers: Arc<Mutex<HashMap<String, BreakpointResponseHandler>>>,
stream_frame_handlers: Arc<Mutex<HashMap<String, BreakpointStreamFrameHandler>>>,
object_response_handler: Arc<Mutex<Option<ObjectResponseHandler>>>,
dead: Arc<std::sync::atomic::AtomicBool>,
_read_thread: Option<std::thread::JoinHandle<()>>,
}
impl BreakpointWebSocketClient {
pub(crate) fn is_dead(&self) -> bool {
self.dead.load(std::sync::atomic::Ordering::Relaxed)
}
pub(crate) fn connect(base_url: &str) -> Result<Self> {
let parsed = Url::parse(base_url)
.map_err(|e| Error::InvalidRequest(format!("bad base URL: {e}")))?;
let ws_scheme = if parsed.scheme() == "https" { "wss" } else { "ws" };
let host = parsed.host_str().unwrap_or("localhost");
let port = parsed.port().unwrap_or(if parsed.scheme() == "https" { 443 } else { 80 });
let path = parsed.path().trim_end_matches('/');
let ws_url = format!("{ws_scheme}://{host}:{port}{path}/_mockserver_callback_websocket");
let (mut socket, _response) = connect(&ws_url)
.map_err(|e| Error::InvalidRequest(format!("WebSocket connect failed: {e}")))?;
let reg_msg = socket
.read()
.map_err(|e| Error::InvalidRequest(format!("WebSocket read registration: {e}")))?;
let reg_text = match reg_msg {
Message::Text(t) => t,
_ => return Err(Error::InvalidRequest("Expected text registration message".into())),
};
let envelope: WsEnvelope = serde_json::from_str(®_text)?;
let client_id = if envelope.msg_type
== "org.mockserver.serialization.model.WebSocketClientIdDTO"
{
let dto: WsClientIdDTO = serde_json::from_str(&envelope.value)?;
dto.client_id
} else {
return Err(Error::InvalidRequest(format!(
"Unexpected registration type: {}",
envelope.msg_type
)));
};
let request_handlers: Arc<Mutex<HashMap<String, BreakpointRequestHandler>>> =
Arc::new(Mutex::new(HashMap::new()));
let response_handlers: Arc<Mutex<HashMap<String, BreakpointResponseHandler>>> =
Arc::new(Mutex::new(HashMap::new()));
let stream_frame_handlers: Arc<Mutex<HashMap<String, BreakpointStreamFrameHandler>>> =
Arc::new(Mutex::new(HashMap::new()));
let object_response_handler: Arc<Mutex<Option<ObjectResponseHandler>>> =
Arc::new(Mutex::new(None));
let socket = Arc::new(Mutex::new(socket));
let dead = Arc::new(std::sync::atomic::AtomicBool::new(false));
let rh = Arc::clone(&request_handlers);
let resh = Arc::clone(&response_handlers);
let sfh = Arc::clone(&stream_frame_handlers);
let orh = Arc::clone(&object_response_handler);
let sock = Arc::clone(&socket);
let dead_flag = Arc::clone(&dead);
let read_thread = std::thread::spawn(move || {
read_loop(sock, rh, resh, sfh, orh, dead_flag);
});
Ok(Self {
client_id,
socket,
request_handlers,
response_handlers,
stream_frame_handlers,
object_response_handler,
dead,
_read_thread: Some(read_thread),
})
}
pub(crate) fn set_request_handler(&self, breakpoint_id: &str, handler: BreakpointRequestHandler) {
self.request_handlers
.lock()
.unwrap()
.insert(breakpoint_id.to_string(), handler);
}
pub(crate) fn set_response_handler(&self, breakpoint_id: &str, handler: BreakpointResponseHandler) {
self.response_handlers
.lock()
.unwrap()
.insert(breakpoint_id.to_string(), handler);
}
pub(crate) fn set_stream_frame_handler(
&self,
breakpoint_id: &str,
handler: BreakpointStreamFrameHandler,
) {
self.stream_frame_handlers
.lock()
.unwrap()
.insert(breakpoint_id.to_string(), handler);
}
pub(crate) fn set_object_response_handler(&self, handler: ObjectResponseHandler) {
*self.object_response_handler.lock().unwrap() = Some(handler);
}
pub(crate) fn remove_handlers(&self, breakpoint_id: &str) {
self.request_handlers.lock().unwrap().remove(breakpoint_id);
self.response_handlers.lock().unwrap().remove(breakpoint_id);
self.stream_frame_handlers
.lock()
.unwrap()
.remove(breakpoint_id);
}
pub(crate) fn clear_handlers(&self) {
self.request_handlers.lock().unwrap().clear();
self.response_handlers.lock().unwrap().clear();
self.stream_frame_handlers.lock().unwrap().clear();
}
pub(crate) fn close(&self) {
if let Ok(mut sock) = self.socket.lock() {
let _ = sock.close(None);
}
}
}
struct ReadLoopDeadGuard(Arc<std::sync::atomic::AtomicBool>);
impl Drop for ReadLoopDeadGuard {
fn drop(&mut self) {
self.0.store(true, std::sync::atomic::Ordering::Relaxed);
}
}
#[allow(clippy::too_many_arguments)]
fn read_loop(
socket: Arc<Mutex<tungstenite::WebSocket<tungstenite::stream::MaybeTlsStream<std::net::TcpStream>>>>,
request_handlers: Arc<Mutex<HashMap<String, BreakpointRequestHandler>>>,
response_handlers: Arc<Mutex<HashMap<String, BreakpointResponseHandler>>>,
stream_frame_handlers: Arc<Mutex<HashMap<String, BreakpointStreamFrameHandler>>>,
object_response_handler: Arc<Mutex<Option<ObjectResponseHandler>>>,
dead: Arc<std::sync::atomic::AtomicBool>,
) {
let _guard = ReadLoopDeadGuard(dead);
loop {
let msg = {
let mut sock = match socket.lock() {
Ok(s) => s,
Err(e) => {
eprintln!("mockserver: breakpoint ws read loop terminated: mutex poisoned: {e}");
return;
}
};
match sock.read() {
Ok(m) => m,
Err(e) => {
eprintln!("mockserver: breakpoint ws read loop terminated: {e}");
return;
}
}
};
let text = match msg {
Message::Text(t) => t,
Message::Close(_) => {
eprintln!("mockserver: breakpoint ws read loop terminated: received close frame");
return;
}
_ => continue,
};
let envelope: WsEnvelope = match serde_json::from_str(&text) {
Ok(e) => e,
Err(_) => continue,
};
match envelope.msg_type.as_str() {
"org.mockserver.model.HttpRequest" => {
handle_request(
&envelope.value,
&request_handlers,
&object_response_handler,
&socket,
);
}
"org.mockserver.model.HttpRequestAndHttpResponse" => {
handle_response(
&envelope.value,
&response_handlers,
&socket,
);
}
"org.mockserver.serialization.model.PausedStreamFrameDTO" => {
handle_stream_frame(
&envelope.value,
&stream_frame_handlers,
&socket,
);
}
"org.mockserver.serialization.model.WebSocketClientIdDTO" => {
}
_ => {}
}
}
}
fn handle_request(
value_json: &str,
handlers: &Arc<Mutex<HashMap<String, BreakpointRequestHandler>>>,
object_response_handler: &Arc<Mutex<Option<ObjectResponseHandler>>>,
socket: &Arc<Mutex<tungstenite::WebSocket<tungstenite::stream::MaybeTlsStream<std::net::TcpStream>>>>,
) {
if let Some((type_name, result)) = route_object_callback(value_json, object_response_handler) {
send_envelope(socket, &type_name, &result);
return;
}
if let Some((type_name, result)) = route_request(value_json, handlers) {
send_envelope(socket, &type_name, &result);
}
}
fn handle_response(
value_json: &str,
handlers: &Arc<Mutex<HashMap<String, BreakpointResponseHandler>>>,
socket: &Arc<Mutex<tungstenite::WebSocket<tungstenite::stream::MaybeTlsStream<std::net::TcpStream>>>>,
) {
if let Some((type_name, result)) = route_response(value_json, handlers) {
send_envelope(socket, &type_name, &result);
}
}
fn handle_stream_frame(
value_json: &str,
handlers: &Arc<Mutex<HashMap<String, BreakpointStreamFrameHandler>>>,
socket: &Arc<Mutex<tungstenite::WebSocket<tungstenite::stream::MaybeTlsStream<std::net::TcpStream>>>>,
) {
if let Some((type_name, decision)) = route_stream_frame(value_json, handlers) {
send_envelope(socket, &type_name, &decision);
}
}
pub fn route_request(
value_json: &str,
handlers: &Arc<Mutex<HashMap<String, BreakpointRequestHandler>>>,
) -> Option<(String, Value)> {
let request: Value = serde_json::from_str(value_json).ok()?;
let correlation_id = extract_header(&request, "WebSocketCorrelationId");
let breakpoint_id = extract_header(&request, "X-MockServer-BreakpointId");
let mut result: Option<Value> = None;
if let Some(bp_id) = &breakpoint_id {
if let Ok(h) = handlers.lock() {
if let Some(handler) = h.get(bp_id) {
result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
handler(request.clone())
}))
.unwrap_or(None);
}
}
}
let mut result = result.unwrap_or_else(|| request.clone());
let type_name = if result.get("statusCode").is_some() {
"org.mockserver.model.HttpResponse"
} else {
"org.mockserver.model.HttpRequest"
};
if let Some(corr) = &correlation_id {
set_header(&mut result, "WebSocketCorrelationId", corr);
}
Some((type_name.to_string(), result))
}
pub fn route_object_callback(
value_json: &str,
handler: &Arc<Mutex<Option<ObjectResponseHandler>>>,
) -> Option<(String, Value)> {
let request: Value = serde_json::from_str(value_json).ok()?;
if extract_header(&request, "X-MockServer-BreakpointId").is_some() {
return None;
}
let correlation_id = extract_header(&request, "WebSocketCorrelationId");
let guard = handler.lock().ok()?;
let handler = guard.as_ref()?;
let mut result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
handler(request.clone())
}))
.unwrap_or_else(|_| serde_json::json!({}));
if let Some(corr) = &correlation_id {
set_header(&mut result, "WebSocketCorrelationId", corr);
}
Some(("org.mockserver.model.HttpResponse".to_string(), result))
}
pub fn route_response(
value_json: &str,
handlers: &Arc<Mutex<HashMap<String, BreakpointResponseHandler>>>,
) -> Option<(String, Value)> {
let req_and_resp: Value = serde_json::from_str(value_json).ok()?;
let http_request = req_and_resp.get("httpRequest").cloned().unwrap_or(Value::Null);
let http_response = req_and_resp.get("httpResponse").cloned().unwrap_or(Value::Null);
let correlation_id = extract_header(&http_request, "WebSocketCorrelationId");
let breakpoint_id = extract_header(&http_request, "X-MockServer-BreakpointId");
let mut result: Option<Value> = None;
if let Some(bp_id) = &breakpoint_id {
if let Ok(h) = handlers.lock() {
if let Some(handler) = h.get(bp_id) {
result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
handler(http_request.clone(), http_response.clone())
}))
.unwrap_or(None);
}
}
}
let mut result = result.unwrap_or(http_response);
if let Some(corr) = &correlation_id {
set_header(&mut result, "WebSocketCorrelationId", corr);
}
Some(("org.mockserver.model.HttpResponse".to_string(), result))
}
pub fn route_stream_frame(
value_json: &str,
handlers: &Arc<Mutex<HashMap<String, BreakpointStreamFrameHandler>>>,
) -> Option<(String, StreamFrameDecision)> {
let frame: PausedStreamFrame = serde_json::from_str(value_json).ok()?;
let mut decision: Option<StreamFrameDecision> = None;
if let Some(bp_id) = &frame.breakpoint_id {
if let Ok(h) = handlers.lock() {
if let Some(handler) = h.get(bp_id) {
decision = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
handler(&frame)
}))
.unwrap_or(None);
}
}
}
let mut decision = decision.unwrap_or_else(|| StreamFrameDecision::continue_frame(&frame.correlation_id));
decision.correlation_id = frame.correlation_id;
Some((
"org.mockserver.serialization.model.StreamFrameDecisionDTO".to_string(),
decision,
))
}
fn send_envelope<T: Serialize>(
socket: &Arc<Mutex<tungstenite::WebSocket<tungstenite::stream::MaybeTlsStream<std::net::TcpStream>>>>,
type_name: &str,
value: &T,
) {
let value_json = match serde_json::to_string(value) {
Ok(j) => j,
Err(_) => return,
};
let envelope = WsEnvelope {
msg_type: type_name.to_string(),
value: value_json,
};
let env_json = match serde_json::to_string(&envelope) {
Ok(j) => j,
Err(_) => return,
};
if let Ok(mut sock) = socket.lock() {
let _ = sock.send(Message::Text(env_json));
}
}
pub fn extract_header(obj: &Value, name: &str) -> Option<String> {
let headers = obj.get("headers")?.as_object()?;
let name_lower = name.to_lowercase();
for (k, v) in headers {
if k.to_lowercase() == name_lower {
if let Some(arr) = v.as_array() {
return arr.first()?.as_str().map(|s| s.to_string());
}
return v.as_str().map(|s| s.to_string());
}
}
None
}
pub fn set_header(obj: &mut Value, name: &str, value: &str) {
if let Some(obj_map) = obj.as_object_mut() {
let headers = obj_map
.entry("headers")
.or_insert_with(|| Value::Object(serde_json::Map::new()));
if let Some(h) = headers.as_object_mut() {
h.insert(
name.to_string(),
serde_json::json!([value]),
);
}
}
}