MetaCrate/crates/libremetaverse-voice-vivox/src/tcp_pipe.rs

//! Cross-platform TCP line transport used by the Vivox gateway adapter.

use crate::{Error, TCPPipeOnDisconnectedCallback, TCPPipeOnReceiveLineCallback};
use libremetaverse_types::compat::{SocketException, Subscription};
use std::collections::HashMap;
use std::io::{Read, Write};
use std::net::{Shutdown, TcpStream};
use std::sync::{
    Arc, Mutex, Weak,
    atomic::{AtomicBool, AtomicU64, Ordering},
};
use std::thread::JoinHandle;
use std::time::Duration;

const READ_POLL: Duration = Duration::from_millis(250);
const MAX_BUFFER_BYTES: usize = 1024 * 1024;

fn lock<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
    mutex
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner)
}

struct Registry<T> {
    next_id: AtomicU64,
    handlers: Arc<Mutex<HashMap<u64, T>>>,
}

impl<T> Default for Registry<T> {
    fn default() -> Self {
        Self {
            next_id: AtomicU64::new(1),
            handlers: Arc::new(Mutex::new(HashMap::new())),
        }
    }
}

impl<T: Send + 'static> Registry<T> {
    fn subscribe(&self, handler: Option<T>) -> Subscription {
        let Some(handler) = handler else {
            return Subscription::detached();
        };
        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
        lock(&self.handlers).insert(id, handler);
        let handlers: Weak<Mutex<HashMap<u64, T>>> = Arc::downgrade(&self.handlers);
        Subscription::new(move || {
            if let Some(handlers) = handlers.upgrade() {
                lock(&handlers).remove(&id);
            }
        })
    }
}

impl<T: Clone> Registry<T> {
    fn snapshot(&self) -> Vec<T> {
        lock(&self.handlers).values().cloned().collect()
    }
}

struct TCPPipeInner {
    writer: Mutex<Option<TcpStream>>,
    reader: Mutex<Option<JoinHandle<()>>>,
    connected: AtomicBool,
    intentional_disconnect: AtomicBool,
    received: Registry<TCPPipeOnReceiveLineCallback>,
    disconnected: Registry<TCPPipeOnDisconnectedCallback>,
}

impl Default for TCPPipeInner {
    fn default() -> Self {
        Self {
            writer: Mutex::new(None),
            reader: Mutex::new(None),
            connected: AtomicBool::new(false),
            intentional_disconnect: AtomicBool::new(false),
            received: Registry::default(),
            disconnected: Registry::default(),
        }
    }
}

pub struct TCPPipe {
    inner: Arc<TCPPipeInner>,
}

impl TCPPipe {
    pub fn new() -> Result<Self, Error> {
        Ok(Self {
            inner: Arc::new(TCPPipeInner::default()),
        })
    }

    #[must_use]
    pub fn subscribe_on_disconnected(
        &self,
        handler: Option<TCPPipeOnDisconnectedCallback>,
    ) -> Subscription {
        self.inner.disconnected.subscribe(handler)
    }

    #[must_use]
    pub fn subscribe_on_receive_line(
        &self,
        handler: Option<TCPPipeOnReceiveLineCallback>,
    ) -> Subscription {
        self.inner.received.subscribe(handler)
    }

    pub fn connect(&self, address: String, port: i32) -> Result<Option<SocketException>, Error> {
        let Ok(port) = u16::try_from(port) else {
            return Ok(Some(SocketException));
        };
        self.disconnect()?;
        let stream = match TcpStream::connect((address.as_str(), port)) {
            Ok(stream) => stream,
            Err(_) => return Ok(Some(SocketException)),
        };
        stream.set_nodelay(true).map_err(|_| Error::Socket)?;
        let reader = stream.try_clone().map_err(|_| Error::Socket)?;
        reader
            .set_read_timeout(Some(READ_POLL))
            .map_err(|_| Error::Socket)?;
        self.inner
            .intentional_disconnect
            .store(false, Ordering::Release);
        self.inner.connected.store(true, Ordering::Release);
        *lock(&self.inner.writer) = Some(stream);
        let inner = Arc::clone(&self.inner);
        *lock(&self.inner.reader) = Some(std::thread::spawn(move || read_loop(inner, reader)));
        Ok(None)
    }

    pub fn disconnect(&self) -> Result<(), Error> {
        self.inner
            .intentional_disconnect
            .store(true, Ordering::Release);
        self.inner.connected.store(false, Ordering::Release);
        if let Some(stream) = lock(&self.inner.writer).take() {
            let _ = stream.shutdown(Shutdown::Both);
        }
        if let Some(reader) = lock(&self.inner.reader).take()
            && reader.thread().id() != std::thread::current().id()
        {
            reader.join().map_err(|_| Error::Socket)?;
        }
        Ok(())
    }

    pub fn send_data(&self, data: Vec<u8>) -> Result<(), Error> {
        if !self.connected() {
            return Err(Error::InvalidOperation);
        }
        let mut writer = lock(&self.inner.writer);
        writer
            .as_mut()
            .ok_or(Error::InvalidOperation)?
            .write_all(&data)
            .map_err(|_| Error::Socket)
    }

    pub fn send_line(&self, message: String) -> Result<(), Error> {
        if message.len() > MAX_BUFFER_BYTES {
            return Err(Error::Argument);
        }
        let mut bytes = Vec::with_capacity(message.len().saturating_add(1));
        bytes.extend(message.chars().map(|character| {
            if character.is_ascii() {
                character as u8
            } else {
                b'?'
            }
        }));
        bytes.push(b'\n');
        self.send_data(bytes)
    }

    #[must_use]
    pub fn connected(&self) -> bool {
        self.inner.connected.load(Ordering::Acquire)
    }
}

impl Drop for TCPPipe {
    fn drop(&mut self) {
        let _ = self.disconnect();
    }
}

fn read_loop(inner: Arc<TCPPipeInner>, mut reader: TcpStream) {
    let mut pending = Vec::new();
    let mut chunk = [0_u8; 4096];
    loop {
        match reader.read(&mut chunk) {
            Ok(0) => break,
            Ok(count) => {
                let data = &chunk[..count];
                let data = data.split(|byte| *byte == 0).next().unwrap_or_default();
                pending.extend_from_slice(data);
                if pending.len() > MAX_BUFFER_BYTES {
                    break;
                }
                drain_lines(&inner, &mut pending);
            }
            Err(error)
                if matches!(
                    error.kind(),
                    std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
                ) =>
            {
                if !inner.connected.load(Ordering::Acquire) {
                    return;
                }
            }
            Err(_) => break,
        }
    }
    let was_connected = inner.connected.swap(false, Ordering::AcqRel);
    lock(&inner.writer).take();
    if was_connected && !inner.intentional_disconnect.load(Ordering::Acquire) {
        for callback in inner.disconnected.snapshot() {
            let _ = callback.invoke(SocketException);
        }
    }
}

fn drain_lines(inner: &TCPPipeInner, pending: &mut Vec<u8>) {
    loop {
        let Some(index) = pending
            .iter()
            .position(|byte| matches!(*byte, b'\r' | b'\n'))
        else {
            return;
        };
        let line = String::from_utf8_lossy(&pending[..index]).into_owned();
        let mut consumed = index + 1;
        if pending.get(index) == Some(&b'\r') && pending.get(consumed) == Some(&b'\n') {
            consumed += 1;
        }
        pending.drain(..consumed);
        if line.trim().is_empty() {
            continue;
        }
        for callback in inner.received.snapshot() {
            let _ = callback.invoke(line.clone());
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::TcpListener;
    use std::sync::mpsc;

    #[test]
    fn connects_sends_and_frames_crlf_and_lf_lines() {
        let listener = TcpListener::bind(("127.0.0.1", 0)).unwrap();
        let port = listener.local_addr().unwrap().port();
        let server = std::thread::spawn(move || {
            let (mut stream, _) = listener.accept().unwrap();
            stream.write_all(b"first\r\nsecond\n").unwrap();
            let mut sent = [0_u8; 6];
            stream.read_exact(&mut sent).unwrap();
            assert_eq!(&sent, b"hello\n");
        });

        let pipe = TCPPipe::new().unwrap();
        let (sender, receiver) = mpsc::channel();
        let _subscription =
            pipe.subscribe_on_receive_line(Some(TCPPipeOnReceiveLineCallback::from_handler(
                move |line| sender.send(line).map_err(|_| Error::Socket),
            )));
        assert!(
            pipe.connect("127.0.0.1".into(), i32::from(port))
                .unwrap()
                .is_none()
        );
        assert!(pipe.connected());
        assert_eq!(
            receiver.recv_timeout(Duration::from_secs(2)).unwrap(),
            "first"
        );
        assert_eq!(
            receiver.recv_timeout(Duration::from_secs(2)).unwrap(),
            "second"
        );
        pipe.send_line("hello".into()).unwrap();
        server.join().unwrap();
        pipe.disconnect().unwrap();
        assert!(!pipe.connected());
    }
}