use crate::stack::{Stack, pop, push};
use crate::value::Value;
use may::net::{TcpListener, TcpStream};
use rustls::{ClientConnection, StreamOwned};
use std::io::{Read, Write};
use std::net::{IpAddr, SocketAddr};
use std::sync::Mutex;
enum StreamKind {
Tcp(TcpStream),
Tls(Box<StreamOwned<ClientConnection, TcpStream>>),
}
impl Read for StreamKind {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
match self {
StreamKind::Tcp(s) => s.read(buf),
StreamKind::Tls(s) => s.read(buf),
}
}
}
impl Write for StreamKind {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
match self {
StreamKind::Tcp(s) => s.write(buf),
StreamKind::Tls(s) => s.write(buf),
}
}
fn flush(&mut self) -> std::io::Result<()> {
match self {
StreamKind::Tcp(s) => s.flush(),
StreamKind::Tls(s) => s.flush(),
}
}
}
const MAX_SOCKETS: usize = 10_000;
#[allow(dead_code)]
const MAX_READ_SIZE: usize = 1_048_576;
struct SocketRegistry<T> {
sockets: Vec<Option<T>>,
free_ids: Vec<usize>,
}
impl<T> SocketRegistry<T> {
const fn new() -> Self {
Self {
sockets: Vec::new(),
free_ids: Vec::new(),
}
}
fn allocate(&mut self, socket: T) -> Result<i64, &'static str> {
if let Some(id) = self.free_ids.pop() {
self.sockets[id] = Some(socket);
return Ok(id as i64);
}
if self.sockets.len() >= MAX_SOCKETS {
return Err("Maximum socket limit reached");
}
let id = self.sockets.len();
self.sockets.push(Some(socket));
Ok(id as i64)
}
fn get_mut(&mut self, id: usize) -> Option<&mut Option<T>> {
self.sockets.get_mut(id)
}
fn free(&mut self, id: usize) {
if let Some(slot) = self.sockets.get_mut(id)
&& slot.is_some()
{
*slot = None;
self.free_ids.push(id);
}
}
fn release_reserved(&mut self, id: usize) {
let is_reserved = self
.sockets
.get(id)
.map(|slot| slot.is_none())
.unwrap_or(false);
if is_reserved && !self.free_ids.contains(&id) {
self.free_ids.push(id);
}
}
}
static LISTENERS: Mutex<SocketRegistry<TcpListener>> = Mutex::new(SocketRegistry::new());
static STREAMS: Mutex<SocketRegistry<StreamKind>> = Mutex::new(SocketRegistry::new());
fn take_tcp(id: usize) -> Option<may::net::TcpStream> {
let mut streams = STREAMS.lock().unwrap();
let slot = streams.get_mut(id)?;
match slot.take() {
Some(StreamKind::Tcp(t)) => Some(t),
Some(other) => {
*slot = Some(other);
None
}
None => None,
}
}
fn release_reserved_stream(id: usize) {
STREAMS.lock().unwrap().release_reserved(id);
}
pub(crate) fn upgrade_tcp_in_place<F>(id: usize, f: F) -> bool
where
F: FnOnce(
may::net::TcpStream,
)
-> Result<rustls::StreamOwned<rustls::ClientConnection, may::net::TcpStream>, ()>,
{
let tcp = match take_tcp(id) {
Some(t) => t,
None => return false,
};
let stream = match f(tcp) {
Ok(s) => s,
Err(()) => {
release_reserved_stream(id);
return false;
}
};
let mut streams = STREAMS.lock().unwrap();
match streams.get_mut(id) {
Some(slot) => {
*slot = Some(StreamKind::Tls(Box::new(stream)));
true
}
None => false,
}
}
const DEFAULT_TCP_CONNECT_TIMEOUT_MS: u64 = 10_000;
static TCP_CONNECT_TIMEOUT: std::sync::LazyLock<std::time::Duration> =
std::sync::LazyLock::new(|| {
let ms = std::env::var("SEQ_TCP_CONNECT_TIMEOUT_MS")
.ok()
.and_then(|v| v.parse::<u64>().ok())
.filter(|n| *n > 0)
.unwrap_or(DEFAULT_TCP_CONNECT_TIMEOUT_MS);
std::time::Duration::from_millis(ms)
});
#[cfg(test)]
static TCP_CONNECT_TIMEOUT_OVERRIDE: Mutex<Option<std::time::Duration>> = Mutex::new(None);
#[cfg(test)]
pub(crate) fn set_test_tcp_connect_timeout(dur: Option<std::time::Duration>) {
*TCP_CONNECT_TIMEOUT_OVERRIDE.lock().unwrap() = dur;
}
fn tcp_connect_timeout() -> std::time::Duration {
#[cfg(test)]
if let Some(dur) = *TCP_CONNECT_TIMEOUT_OVERRIDE.lock().unwrap() {
return dur;
}
*TCP_CONNECT_TIMEOUT
}
pub(crate) fn connect_to_addrs(addrs: &[IpAddr], port: u16) -> Option<TcpStream> {
let timeout = tcp_connect_timeout();
addrs
.iter()
.find_map(|ip| TcpStream::connect_timeout(&SocketAddr::new(*ip, port), timeout).ok())
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn patch_seq_tcp_listen(stack: Stack) -> Stack {
unsafe {
let (stack, port_val) = pop(stack);
let port = match port_val {
Value::Int(p) => p,
_ => {
let stack = push(stack, Value::Int(0));
return push(stack, Value::Bool(false));
}
};
if !(0..=65535).contains(&port) {
let stack = push(stack, Value::Int(0));
return push(stack, Value::Bool(false));
}
let addr = format!("0.0.0.0:{}", port);
let listener = match TcpListener::bind(&addr) {
Ok(l) => l,
Err(_) => {
let stack = push(stack, Value::Int(0));
return push(stack, Value::Bool(false));
}
};
let mut listeners = LISTENERS.lock().unwrap();
match listeners.allocate(listener) {
Ok(listener_id) => {
let stack = push(stack, Value::Int(listener_id));
push(stack, Value::Bool(true))
}
Err(_) => {
let stack = push(stack, Value::Int(0));
push(stack, Value::Bool(false))
}
}
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn patch_seq_tcp_connect(stack: Stack) -> Stack {
unsafe {
let (stack, port_val) = pop(stack);
let port = match port_val {
Value::Int(p) => p,
_ => {
let stack = push(stack, Value::Int(0));
return push(stack, Value::Bool(false));
}
};
if !(1..=65535).contains(&port) {
let stack = push(stack, Value::Int(0));
return push(stack, Value::Bool(false));
}
let (stack, host_val) = pop(stack);
let host = match host_val {
Value::String(s) => s,
_ => {
let stack = push(stack, Value::Int(0));
return push(stack, Value::Bool(false));
}
};
let hostname = host.as_str_or_empty();
if hostname.is_empty() {
let stack = push(stack, Value::Int(0));
return push(stack, Value::Bool(false));
}
let addrs = crate::dns::resolve_to_ips(hostname);
if addrs.is_empty() {
let stack = push(stack, Value::Int(0));
return push(stack, Value::Bool(false));
}
let stream = match connect_to_addrs(&addrs, port as u16) {
Some(s) => s,
None => {
let stack = push(stack, Value::Int(0));
return push(stack, Value::Bool(false));
}
};
let mut streams = STREAMS.lock().unwrap();
match streams.allocate(StreamKind::Tcp(stream)) {
Ok(id) => {
let stack = push(stack, Value::Int(id));
push(stack, Value::Bool(true))
}
Err(_) => {
let stack = push(stack, Value::Int(0));
push(stack, Value::Bool(false))
}
}
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn patch_seq_tcp_accept(stack: Stack) -> Stack {
unsafe {
let (stack, listener_id_val) = pop(stack);
let listener_id = match listener_id_val {
Value::Int(id) => id as usize,
_ => {
let stack = push(stack, Value::Int(0));
return push(stack, Value::Bool(false));
}
};
let listener = {
let mut listeners = LISTENERS.lock().unwrap();
match listeners.get_mut(listener_id).and_then(|opt| opt.take()) {
Some(l) => l,
None => {
let stack = push(stack, Value::Int(0));
return push(stack, Value::Bool(false));
}
}
};
let (stream, _addr) = match listener.accept() {
Ok(result) => result,
Err(_) => {
let mut listeners = LISTENERS.lock().unwrap();
if let Some(slot) = listeners.get_mut(listener_id) {
*slot = Some(listener);
}
let stack = push(stack, Value::Int(0));
return push(stack, Value::Bool(false));
}
};
{
let mut listeners = LISTENERS.lock().unwrap();
if let Some(slot) = listeners.get_mut(listener_id) {
*slot = Some(listener);
}
}
let mut streams = STREAMS.lock().unwrap();
match streams.allocate(StreamKind::Tcp(stream)) {
Ok(client_id) => {
let stack = push(stack, Value::Int(client_id));
push(stack, Value::Bool(true))
}
Err(_) => {
let stack = push(stack, Value::Int(0));
push(stack, Value::Bool(false))
}
}
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn patch_seq_tcp_read(stack: Stack) -> Stack {
unsafe {
let (stack, socket_id_val) = pop(stack);
let socket_id = match socket_id_val {
Value::Int(id) => id as usize,
_ => {
let stack = push(stack, Value::String("".into()));
return push(stack, Value::Bool(false));
}
};
let mut stream = {
let mut streams = STREAMS.lock().unwrap();
match streams.get_mut(socket_id).and_then(|opt| opt.take()) {
Some(s) => s,
None => {
let stack = push(stack, Value::String("".into()));
return push(stack, Value::Bool(false));
}
}
};
let mut buffer = Vec::new();
let mut chunk = [0u8; 4096];
let mut read_error = false;
match stream.read(&mut chunk) {
Ok(0) => {} Ok(n) => buffer.extend_from_slice(&chunk[..n]),
Err(_) => read_error = true,
}
{
let mut streams = STREAMS.lock().unwrap();
if let Some(slot) = streams.get_mut(socket_id) {
*slot = Some(stream);
}
}
if read_error {
let stack = push(stack, Value::String("".into()));
return push(stack, Value::Bool(false));
}
let stack = push(stack, Value::String(crate::seqstring::global_bytes(buffer)));
push(stack, Value::Bool(true))
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn patch_seq_tcp_write(stack: Stack) -> Stack {
unsafe {
let (stack, socket_id_val) = pop(stack);
let socket_id = match socket_id_val {
Value::Int(id) => id as usize,
_ => {
return push(stack, Value::Bool(false));
}
};
let (stack, data_val) = pop(stack);
let data = match data_val {
Value::String(s) => s,
_ => {
return push(stack, Value::Bool(false));
}
};
let mut stream = {
let mut streams = STREAMS.lock().unwrap();
match streams.get_mut(socket_id).and_then(|opt| opt.take()) {
Some(s) => s,
None => {
return push(stack, Value::Bool(false));
}
}
};
let write_result = stream.write_all(data.as_bytes());
let flush_result = if write_result.is_ok() {
stream.flush()
} else {
write_result
};
{
let mut streams = STREAMS.lock().unwrap();
if let Some(slot) = streams.get_mut(socket_id) {
*slot = Some(stream);
}
}
push(stack, Value::Bool(flush_result.is_ok()))
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn patch_seq_tcp_close(stack: Stack) -> Stack {
unsafe {
let (stack, socket_id_val) = pop(stack);
let socket_id = match socket_id_val {
Value::Int(id) => id as usize,
_ => {
return push(stack, Value::Bool(false));
}
};
{
let mut streams = STREAMS.lock().unwrap();
if streams
.get_mut(socket_id)
.is_some_and(|slot| slot.is_some())
{
streams.free(socket_id);
return push(stack, Value::Bool(true));
}
}
{
let mut listeners = LISTENERS.lock().unwrap();
if listeners
.get_mut(socket_id)
.is_some_and(|slot| slot.is_some())
{
listeners.free(socket_id);
return push(stack, Value::Bool(true));
}
}
push(stack, Value::Bool(false))
}
}
pub use patch_seq_tcp_accept as tcp_accept;
pub use patch_seq_tcp_close as tcp_close;
pub use patch_seq_tcp_connect as tcp_connect;
pub use patch_seq_tcp_listen as tcp_listen;
pub use patch_seq_tcp_local_port as tcp_local_port;
pub use patch_seq_tcp_read as tcp_read;
pub use patch_seq_tcp_write as tcp_write;
#[unsafe(no_mangle)]
pub unsafe extern "C" fn patch_seq_tcp_local_port(stack: Stack) -> Stack {
unsafe {
let (stack, socket_id_val) = pop(stack);
let socket_id = match socket_id_val {
Value::Int(id) => id as usize,
_ => {
let stack = push(stack, Value::Int(0));
return push(stack, Value::Bool(false));
}
};
let port: Option<u16> = {
let mut streams = STREAMS.lock().unwrap();
streams
.get_mut(socket_id)
.and_then(|slot| slot.as_ref())
.and_then(|sk| match sk {
StreamKind::Tcp(s) => s.local_addr().ok().map(|a| a.port()),
StreamKind::Tls(s) => s.sock.local_addr().ok().map(|a| a.port()),
})
};
if let Some(port) = port {
let stack = push(stack, Value::Int(port as i64));
return push(stack, Value::Bool(true));
}
let port: Option<u16> = {
let mut listeners = LISTENERS.lock().unwrap();
listeners
.get_mut(socket_id)
.and_then(|slot| slot.as_ref())
.and_then(|l| l.local_addr().ok())
.map(|a| a.port())
};
if let Some(port) = port {
let stack = push(stack, Value::Int(port as i64));
return push(stack, Value::Bool(true));
}
let stack = push(stack, Value::Int(0));
push(stack, Value::Bool(false))
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn patch_seq_socket_cast(stack: Stack) -> Stack {
assert!(!stack.is_null(), "fd<->socket cast: stack is empty");
let (rest, val) = unsafe { pop(stack) };
match val {
Value::Int(fd) => unsafe { push(rest, Value::Int(fd)) },
_ => panic!("fd<->socket cast: expected Int on stack"),
}
}
#[cfg(test)]
mod tests;