use crate::host::{invoke, with_host, JsObj};
use fusevm::Value;
use indexmap::IndexMap;
use std::collections::HashMap;
use std::net::UdpSocket;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;
pub const MODULE_METHODS: &[&str] = &["createSocket"];
pub const SOCKET_TAG: &str = "UdpSocket";
pub const SOCKET_METHODS: &[&str] = &[
"bind",
"send",
"close",
"address",
"setBroadcast",
"setTTL",
"setMulticastTTL",
"setMulticastLoopback",
"setMulticastInterface",
"addMembership",
"dropMembership",
"addSourceSpecificMembership",
"dropSourceSpecificMembership",
"setRecvBufferSize",
"setSendBufferSize",
"getRecvBufferSize",
"getSendBufferSize",
"connect",
"disconnect",
"remoteAddress",
"ref",
"unref",
];
const POLL: Duration = Duration::from_millis(200);
struct UdpRec {
emitter: Value,
socket: Arc<UdpSocket>,
stop: Arc<AtomicBool>,
}
#[derive(Default)]
struct DgramState {
next_id: u64,
sockets: HashMap<u64, UdpRec>,
}
thread_local! {
static DGRAM: std::cell::RefCell<DgramState> = std::cell::RefCell::new(DgramState::default());
}
fn next_id() -> u64 {
DGRAM.with(|s| {
let mut s = s.borrow_mut();
s.next_id += 1;
s.next_id
})
}
fn get_prop(recv: &Value, key: &str) -> Option<Value> {
with_host(|h| match h.get(recv) {
Some(JsObj::Object(p)) => p.get(key).cloned(),
_ => None,
})
}
fn set_prop(recv: &Value, key: &str, val: Value) {
with_host(|h| {
if let Some(JsObj::Object(p)) = h.get_mut(recv) {
p.insert(key.to_string(), val);
}
});
}
fn u64_prop(recv: &Value, key: &str) -> Option<u64> {
get_prop(recv, key).map(|v| with_host(|h| h.to_number(&v)) as u64)
}
fn is_udp6(recv: &Value) -> bool {
get_prop(recv, "@@udptype")
.map(|v| with_host(|h| h.str_of(&v)))
.as_deref()
== Some("udp6")
}
fn default_bind_host(recv: &Value) -> &'static str {
if is_udp6(recv) {
"::"
} else {
"0.0.0.0"
}
}
fn default_send_host(recv: &Value) -> &'static str {
if is_udp6(recv) {
"::1"
} else {
"127.0.0.1"
}
}
fn is_num(v: &Value) -> bool {
matches!(v, Value::Float(_) | Value::Int(_))
}
fn is_str(v: &Value) -> bool {
matches!(v, Value::Str(_)) || with_host(|h| matches!(h.get(v), Some(JsObj::Str(_))))
}
fn value_bytes(v: &Value) -> Vec<u8> {
let is_buffer =
with_host(|h| matches!(h.get(v), Some(JsObj::Object(p)) if p.contains_key("@@bytes")));
if is_buffer {
return with_host(|h| match h.get(v) {
Some(JsObj::Object(p)) => match p.get("@@bytes").and_then(|b| h.get(b)) {
Some(JsObj::Array(items)) => items.iter().map(|x| h.to_number(x) as u8).collect(),
_ => Vec::new(),
},
_ => Vec::new(),
});
}
with_host(|h| h.str_of(v)).into_bytes()
}
fn emitter_dispatch(recv: &Value, method: &str, args: &[Value]) -> Option<Result<Value, String>> {
match method {
"on"
| "addListener"
| "prependListener"
| "once"
| "prependOnceListener"
| "emit"
| "removeListener"
| "off"
| "removeAllListeners"
| "listeners"
| "listenerCount"
| "eventNames"
| "setMaxListeners"
| "getMaxListeners" => Some(super::events::instance_call(recv, method, args.to_vec())),
_ => None,
}
}
pub fn call(method: &str, args: &[Value]) -> Option<Result<Value, String>> {
match method {
"createSocket" => Some(Ok(create_socket(args))),
_ => None,
}
}
pub fn create_socket(args: &[Value]) -> Value {
let first = args.first().cloned().unwrap_or(Value::Undef);
let sock_type = if is_str(&first) {
with_host(|h| h.str_of(&first))
} else {
with_host(|h| match h.get(&first) {
Some(JsObj::Object(p)) => p.get("type").map(|v| h.str_of(v)),
_ => None,
})
.unwrap_or_else(|| "udp4".to_string())
};
let sock_type = if sock_type == "udp6" { "udp6" } else { "udp4" };
let mut extra = IndexMap::new();
extra.insert("@@udptype".into(), with_host(|h| h.new_str(sock_type)));
let socket = super::net::new_emitter_object(SOCKET_TAG, extra);
if let Some(cb) = args
.get(1)
.filter(|v| with_host(|h| crate::host::is_callable(h, v)))
{
let _ = super::events::instance_call(
&socket,
"on",
vec![with_host(|h| h.new_str("message")), cb.clone()],
);
}
socket
}
pub fn instance_call(recv: &Value, method: &str, args: Vec<Value>) -> Result<Value, String> {
if let Some(r) = emitter_dispatch(recv, method, &args) {
return r;
}
match method {
"bind" => socket_bind(recv, &args),
"send" => socket_send(recv, &args),
"close" => socket_close(recv, &args),
"address" => socket_address(recv),
"setBroadcast" => {
let on = with_host(|h| h.truthy(args.first().unwrap_or(&Value::Undef)));
if let Some(sock) = live_socket(recv) {
sock.set_broadcast(on).ok();
}
Ok(recv.clone())
}
"setTTL" | "setMulticastTTL" => {
let ttl = with_host(|h| h.to_number(args.first().unwrap_or(&Value::Undef))) as u32;
if let Some(sock) = live_socket(recv) {
if method == "setTTL" {
sock.set_ttl(ttl).ok();
} else {
sock.set_multicast_ttl_v4(ttl).ok();
}
}
Ok(args.first().cloned().unwrap_or(Value::Undef))
}
"getRecvBufferSize" | "getSendBufferSize" => Ok(Value::Float(65536.0)),
"setMulticastLoopback"
| "setMulticastInterface"
| "addMembership"
| "dropMembership"
| "addSourceSpecificMembership"
| "dropSourceSpecificMembership"
| "setRecvBufferSize"
| "setSendBufferSize"
| "connect"
| "disconnect"
| "remoteAddress"
| "ref"
| "unref" => Ok(recv.clone()),
_ => Err(crate::host::type_error(&format!(
"socket.{method} is not a function"
))),
}
}
fn live_socket(recv: &Value) -> Option<Arc<UdpSocket>> {
let id = u64_prop(recv, "@@dgramid")?;
DGRAM.with(|s| s.borrow().sockets.get(&id).map(|r| r.socket.clone()))
}
fn socket_bind(recv: &Value, args: &[Value]) -> Result<Value, String> {
let mut port: u16 = 0;
let mut host = default_bind_host(recv).to_string();
let mut cb: Option<Value> = None;
if let Some(first) = args.first() {
if is_num(first) {
port = with_host(|h| h.to_number(first)) as u16;
} else if with_host(
|h| matches!(h.get(first), Some(JsObj::Object(p)) if !p.contains_key("@@native")),
) {
with_host(|h| {
if let Some(JsObj::Object(p)) = h.get(first) {
if let Some(pv) = p.get("port") {
port = h.to_number(pv) as u16;
}
if let Some(av) = p.get("address").map(|v| h.str_of(v)) {
host = av;
}
}
});
}
}
for a in args.iter().skip(1) {
if is_str(a) {
host = with_host(|h| h.str_of(a));
} else if with_host(|h| crate::host::is_callable(h, a)) {
cb = Some(a.clone());
}
}
do_bind(recv, &host, port)?;
let socket = recv.clone();
let _ = with_host(|h| h.io_sender()).send(Box::new(move || {
if let Some(cb) = cb {
super::events::instance_call(
&socket,
"once",
vec![with_host(|h| h.new_str("listening")), cb],
)?;
}
super::events::instance_call(&socket, "emit", vec![with_host(|h| h.new_str("listening"))])?;
Ok(())
}));
Ok(recv.clone())
}
fn do_bind(recv: &Value, host: &str, port: u16) -> Result<Arc<UdpSocket>, String> {
if let Some(sock) = live_socket(recv) {
return Ok(sock);
}
let socket =
UdpSocket::bind((host, port)).map_err(|e| format!("Error: bind EADDRINUSE: {e}"))?;
socket.set_read_timeout(Some(POLL)).ok();
let socket = Arc::new(socket);
let id = next_id();
set_prop(recv, "@@dgramid", Value::Float(id as f64));
let stop = Arc::new(AtomicBool::new(false));
DGRAM.with(|s| {
s.borrow_mut().sockets.insert(
id,
UdpRec {
emitter: recv.clone(),
socket: socket.clone(),
stop: stop.clone(),
},
);
});
with_host(|h| h.incr_handle());
let tx = with_host(|h| h.io_sender());
let recv_sock = socket.clone();
std::thread::spawn(move || recv_loop(recv_sock, id, stop, tx));
Ok(socket)
}
fn recv_loop(
socket: Arc<UdpSocket>,
id: u64,
stop: Arc<AtomicBool>,
tx: std::sync::mpsc::Sender<crate::host::IoTask>,
) {
let mut buf = [0u8; 65536];
loop {
if stop.load(Ordering::Acquire) {
break;
}
match socket.recv_from(&mut buf) {
Ok((n, src)) => {
let bytes = buf[..n].to_vec();
let address = src.ip().to_string();
let port = src.port();
let family = if src.is_ipv6() { "IPv6" } else { "IPv4" };
let _ = tx.send(Box::new(move || {
on_message(id, bytes, address, port, family)
}));
}
Err(ref e)
if e.kind() == std::io::ErrorKind::WouldBlock
|| e.kind() == std::io::ErrorKind::TimedOut =>
{
continue;
}
Err(_) => break,
}
}
}
fn on_message(
id: u64,
bytes: Vec<u8>,
address: String,
port: u16,
family: &'static str,
) -> Result<(), String> {
let socket = DGRAM.with(|s| s.borrow().sockets.get(&id).map(|r| r.emitter.clone()));
let Some(socket) = socket else { return Ok(()) };
let size = bytes.len();
let msg = super::buffer::from_bytes(&bytes);
let rinfo = with_host(|h| {
let mut m = IndexMap::new();
m.insert("address".into(), h.new_str(address));
m.insert("family".into(), h.new_str(family));
m.insert("port".into(), Value::Float(port as f64));
m.insert("size".into(), Value::Float(size as f64));
h.new_object(m)
});
super::events::instance_call(
&socket,
"emit",
vec![with_host(|h| h.new_str("message")), msg, rinfo],
)?;
Ok(())
}
fn socket_send(recv: &Value, args: &[Value]) -> Result<Value, String> {
let msg = args.first().cloned().unwrap_or(Value::Undef);
let full = value_bytes(&msg);
let mut nums: Vec<f64> = Vec::new();
let mut i = 1;
while i < args.len() && is_num(&args[i]) {
nums.push(with_host(|h| h.to_number(&args[i])));
i += 1;
}
let (offset, length, port) = if nums.len() >= 3 {
(
nums[0].max(0.0) as usize,
nums[1].max(0.0) as usize,
nums[2] as u16,
)
} else if let Some(p) = nums.first() {
(0usize, full.len(), *p as u16)
} else {
return Err(crate::host::type_error("Port should be > 0 and < 65536"));
};
let mut address = default_send_host(recv).to_string();
let mut cb: Option<Value> = None;
for a in args.iter().skip(i) {
if is_str(a) {
address = with_host(|h| h.str_of(a));
} else if with_host(|h| crate::host::is_callable(h, a)) {
cb = Some(a.clone());
}
}
let end = offset.saturating_add(length).min(full.len());
let start = offset.min(full.len());
let data = &full[start..end.max(start)];
let socket = do_bind(recv, default_bind_host(recv), 0)?;
socket
.send_to(data, (address.as_str(), port))
.map_err(|e| format!("Error: send {e}"))?;
if let Some(cb) = cb {
let _ = with_host(|h| h.io_sender()).send(Box::new(move || {
let nul = with_host(|h| h.null());
invoke(&cb, vec![nul], None)?;
Ok(())
}));
}
Ok(Value::Undef)
}
fn socket_close(recv: &Value, args: &[Value]) -> Result<Value, String> {
if let Some(id) = u64_prop(recv, "@@dgramid") {
let rec = DGRAM.with(|s| s.borrow_mut().sockets.remove(&id));
if let Some(rec) = rec {
rec.stop.store(true, Ordering::Release);
with_host(|h| h.decr_handle());
let _ = with_host(|h| h.io_sender()).send(Box::new(|| Ok(())));
}
}
if let Some(cb) = args
.first()
.filter(|v| with_host(|h| crate::host::is_callable(h, v)))
{
invoke(cb, Vec::new(), None)?;
}
super::events::instance_call(recv, "emit", vec![with_host(|h| h.new_str("close"))])?;
Ok(Value::Undef)
}
fn socket_address(recv: &Value) -> Result<Value, String> {
let socket = live_socket(recv)
.ok_or_else(|| "Error: getsockname EBADF: bad file descriptor".to_string())?;
let addr = socket
.local_addr()
.map_err(|e| format!("Error: getsockname {e}"))?;
Ok(with_host(|h| {
let mut m = IndexMap::new();
m.insert("address".into(), h.new_str(addr.ip().to_string()));
m.insert(
"family".into(),
h.new_str(if addr.is_ipv6() { "IPv6" } else { "IPv4" }),
);
m.insert("port".into(), Value::Float(addr.port() as f64));
h.new_object(m)
}))
}