Make streams system

2026-06-08 22:38:01 -06:00 · 2025-11-25 17:31:09 -07:00
parent 6863e08a0a
commit b43f2f5181
10 changed files with 395 additions and 202 deletions
@@ -0,0 +1,126 @@
+use std::{
+    io::{Read, Write},
+    net,
+    sync::{
+        Arc,
+        atomic::{AtomicBool, Ordering},
+    },
+};
+
+use crate::{Announcement, ModuleError, debug, network::Stream};
+
+pub struct TcpStream(Arc<AtomicBool>, net::TcpStream);
+
+impl TcpStream {
+    pub fn new(stream: net::TcpStream) -> Self {
+        stream.set_nonblocking(true).unwrap();
+        Self(Arc::new(AtomicBool::new(true)), stream)
+    }
+
+    // Call this when the stream ends
+    fn disconnected(&mut self) {
+        self.0.store(false, Ordering::Relaxed);
+    }
+}
+
+impl Stream<Announcement> for TcpStream {
+    fn is_alive(&self) -> bool {
+        // if self.1.take_error().unwrap_or(None).is_some() {
+        //     // self.1.pe
+        //     warn!("Disconnected #################");
+        //     return true;
+        // } else {
+        //     return false;
+        // }
+
+        // let mut buf = [0u8; 1];
+        // match self.1.peek(&mut buf) {
+        //     Ok(n) => n == 1,
+        //     Err(_) => false,
+        // }
+
+        let mut buf = [0u8; 1];
+        match self.1.peek(&mut buf) {
+            Ok(0) => false, // Connection closed (EOF)
+            Ok(_) => true,  // Data available or connection alive
+            Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => true, // No data but alive
+            Err(_) => false, // Connection error
+        }
+
+        // true
+
+        // self.0.load(Ordering::Relaxed)
+    }
+
+    fn has_recv(&self) -> bool {
+        let mut buf = [0u8; 1];
+        match self.1.peek(&mut buf) {
+            Ok(n) if n > 0 => true, // Data is available
+            Ok(_) => false,         // EOF (connection closed)
+            Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => false, // No data
+            Err(_) => false,
+        }
+        // false
+    }
+
+    fn read(&mut self) -> Vec<Announcement> {
+        let mut ret = Vec::new();
+
+        while self.has_recv() {
+            let mut size_buf = [0u8; 4];
+            match self.1.read_exact(&mut size_buf) {
+                Ok(()) => {}
+                Err(_) => {
+                    self.disconnected();
+                    break;
+                }
+            };
+            let size = u32::from_be_bytes(size_buf);
+
+            let mut buf = vec![0u8; size as usize];
+
+            match self.1.read_exact(&mut buf) {
+                Ok(()) => {}
+                Err(_) => {
+                    self.disconnected();
+                }
+            }
+
+            if let Some(a) = Announcement::decode(&buf) {
+                ret.push(a);
+            } else {
+                debug!("Malformed data");
+            }
+        }
+
+        ret
+    }
+
+    fn write(&mut self, announcement: Announcement) -> Result<(), crate::ModuleError> {
+        let bytes = announcement.encode();
+
+        // Write length of bytes
+        self.1
+            .write_all(&u32::to_be_bytes(bytes.len() as u32))
+            .map_err(|e| ModuleError::Error(e.to_string().into()))?;
+        // Write data
+        self.1
+            .write_all(&bytes)
+            .map_err(|e| ModuleError::Error(e.to_string().into()))?;
+        // Flush data
+        self.1
+            .flush()
+            .map_err(|e| ModuleError::Error(e.to_string().into()))?;
+
+        Ok(())
+    }
+
+    fn try_clone(&self) -> Result<Box<dyn Stream<Announcement> + Send + Sync>, crate::ModuleError> {
+        Ok(Box::new(Self(
+            self.0.clone(),
+            self.1
+                .try_clone()
+                .map_err(|e| ModuleError::Error(e.to_string().into()))?,
+        )))
+    }
+}