worldwideportal/src/websocket.rs

169 lines
6.1 KiB
Rust

use std::collections::BTreeMap;
use gc_arena::Collect;
use serde::Deserialize;
use wasm_bindgen::{closure::Closure, JsCast, JsValue};
use wasm_bindgen_futures::js_sys::Date;
use web_sys::{
console,
js_sys::{ArrayBuffer, JsString, Uint8Array},
BinaryType, DomException, MessageEvent, WebSocket,
};
use crate::{
logging::{self, LogEvent},
lua_engine::muds::{handle_websocket_close_log_err, handle_websocket_output_log_err},
GlobalMemoCell,
};
#[derive(PartialEq, PartialOrd, Eq, Ord, Debug, Clone, Collect)]
#[collect(require_static)]
pub struct WebSocketId(pub u64);
pub struct WebSocketData {
pub connection: WebSocket,
pub trusted: bool, // Is it allowed to submit Lua to be run?
pub closed: bool,
pub url: String,
pub retained_closures: Option<(Closure<dyn FnMut(MessageEvent)>, Closure<dyn FnMut()>)>,
pub log_dest: Option<String>,
}
pub type RegisteredWebSockets = BTreeMap<WebSocketId, WebSocketData>;
#[derive(Deserialize)]
pub enum MessageFromServer {
RunLua(String),
#[serde(other)]
Unknown,
}
pub fn connect_websocket(
trusted: bool,
url: &str,
sockets: &mut RegisteredWebSockets,
// Only used for callbacks, expected sockets already borrowed mut!
globals: &GlobalMemoCell,
) -> anyhow::Result<WebSocketId> {
let mut data = WebSocketData {
connection: WebSocket::new(url).map_err(|e| {
anyhow::Error::msg(
e.dyn_into::<DomException>()
.map(|e| e.message().to_string())
.unwrap_or("Unknown error connecting".to_owned()),
)
})?,
trusted,
closed: false,
url: url.to_owned(),
retained_closures: None,
log_dest: None,
};
data.connection.set_binary_type(BinaryType::Arraybuffer);
let new_id = sockets
.last_key_value()
.map_or(WebSocketId(0), |v| WebSocketId(v.0 .0 + 1));
let data_globals = globals.clone();
let data_new_id = new_id.clone();
let data_closure: Closure<dyn FnMut(MessageEvent)> = Closure::new(move |ev: MessageEvent| {
let data = ev.data();
if data.has_type::<ArrayBuffer>() {
let log_dest = data_globals
.ws_registry
.borrow()
.get(&data_new_id)
.and_then(|d| d.log_dest.clone());
if let Some(log_dest) = &log_dest {
let str_rendering =
String::from_utf8_lossy(&Uint8Array::new(&data).to_vec()).to_string();
logging::log(
&data_globals,
&LogEvent {
log_stream: log_dest.clone(),
timestamp: Date::new_0(),
message: format!("R {}", &str_rendering),
},
);
}
handle_websocket_output_log_err(
&data_new_id,
&data_globals,
&Uint8Array::new(&data).to_vec(),
);
} else if data.has_type::<JsString>() {
let msg: String = data.unchecked_into::<JsString>().into();
if let Some(latest_data) = data_globals.ws_registry.borrow().get(&data_new_id) {
if latest_data.trusted {
match data_globals.lua_engine.borrow_mut().execute(&msg) {
Ok(()) => {}
Err(e) => console::log_3(
&JsValue::from_str("Error executing Lua from websocket"),
&JsValue::from_str(&latest_data.url),
&JsValue::from_str(&e),
),
}
} else {
console::log_2(
&JsValue::from_str("Ignoring Lua from untrusted websocket."),
&JsValue::from_str(&latest_data.url),
);
}
}
}
});
let close_globals = globals.clone();
let close_id = new_id.clone();
let close_closure: Closure<dyn FnMut()> =
Closure::new(
move || match close_globals.ws_registry.borrow_mut().get_mut(&close_id) {
None => {}
Some(closed_socket) => {
closed_socket.closed = true;
closed_socket.retained_closures = None;
handle_websocket_close_log_err(&close_id, &close_globals);
}
},
);
data.connection
.add_event_listener_with_callback("message", data_closure.as_ref().unchecked_ref())
.expect("Couldn't set message handler on WebSocket");
data.connection
.add_event_listener_with_callback("close", close_closure.as_ref().unchecked_ref())
.expect("Couldn't set close handler on WebSocket");
data.connection
.add_event_listener_with_callback("error", close_closure.as_ref().unchecked_ref())
.expect("Couldn't set error handler on WebSocket");
data.retained_closures = Some((data_closure, close_closure));
sockets.insert(new_id.clone(), data);
Ok(new_id)
}
pub fn send_message_to_mud(
socket: &WebSocketId,
msg: &[u8],
global: &GlobalMemoCell,
) -> anyhow::Result<()> {
match global.ws_registry.borrow().get(socket) {
None => Err(anyhow::Error::msg("MUD connection not found")),
Some(sock_data) if sock_data.closed => Err(anyhow::Error::msg("MUD connection is closed")),
Some(sock_data) => {
if let Some(log_dest) = &sock_data.log_dest {
logging::log(
global,
&LogEvent {
log_stream: log_dest.clone(),
timestamp: Date::new_0(),
message: format!("W {}", String::from_utf8_lossy(msg)),
},
);
}
sock_data.connection.send_with_u8_array(msg).map_err(|e| {
e.dyn_into::<DomException>()
.map(|e| anyhow::Error::msg(e.message()))
.unwrap_or_else(|_| anyhow::Error::msg("Unexpected exception while sending"))
})?;
Ok(())
}
}
}