From 932485e5ab4c017978338c80e380a154371425e1 Mon Sep 17 00:00:00 2001
From: pjht <pterpstra@pterpstra.com>
Date: Fri, 20 Sep 2024 14:42:25 -0500
Subject: [PATCH] Add poll support

---
 Cargo.lock  |  4 ++++
 src/main.rs | 65 ++++++++++++++++++++++++++++++++++++++++++++++++-----
 2 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 0a8fcbb..25de524 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -22,6 +22,9 @@ name = "bitflags"
 version = "2.6.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b048fb63fd8b5923fc5aa7b340d8e156aec7ec02f0c78fa8a6ddc2613f6f71de"
+dependencies = [
+ "serde",
+]
 
 [[package]]
 name = "byteorder"
@@ -81,6 +84,7 @@ checksum = "edd0f118536f44f5ccd48bcb8b111bdc3de888b58c74639dfb034a357d0f206d"
 name = "file_rpc"
 version = "0.1.0"
 dependencies = [
+ "bitflags",
  "parking_lot",
  "postcard",
  "serde",
diff --git a/src/main.rs b/src/main.rs
index 5c01215..8855c35 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,16 +1,13 @@
 mod termios;
 
 use std::{
-    borrow::Cow,
-    collections::HashMap,
-    os::mikros::{ipc, syscalls},
-    path::{Path, PathBuf},
-    sync::{
+    borrow::Cow, collections::HashMap, io::empty, os::mikros::{ipc, syscalls}, path::{Path, PathBuf}, sync::{
         atomic::{AtomicBool, AtomicUsize, Ordering},
         Arc,
-    },
+    }
 };
 
+use file_rpc::PollEvents;
 use parking_lot::{Mutex, RwLock};
 use slab::Slab;
 use termios::Termios;
@@ -189,6 +186,22 @@ impl Pty {
             self.column.fetch_sub(1, Ordering::Relaxed);
         }
     }
+
+    fn curr_slave_poll(&self) -> PollEvents {
+        let mut events = PollEvents::POLLOUT;
+        if self.input_buffer.lock().contains(&b'\n') {
+            events |= PollEvents::POLLIN;
+        }
+        events
+    }
+
+    fn curr_master_poll(&self) -> PollEvents {
+        let mut events = PollEvents::POLLOUT;
+        if !self.output_buffer.lock().is_empty() {
+            events |= PollEvents::POLLOUT;
+        }
+        events
+    }
 }
 
 #[derive(Clone)]
@@ -196,6 +209,7 @@ struct Serv {
     ptys: Arc<RwLock<Slab<Pty>>>,
     slave_map: Arc<RwLock<HashMap<PathBuf, usize>>>,
     devfs_client: devfs_rpc::Client,
+    polls: Arc<Mutex<Slab<(u64, PollEvents)>>>,
 }
 
 impl dev_driver_rpc::Server for Serv {
@@ -254,6 +268,44 @@ impl file_rpc::Server for Serv {
     fn dup(&self, fd: u64) -> Option<u64> {
         Some(fd)
     }
+
+    fn register_poll(&self, fd: u64, events: PollEvents) -> u64 {
+        self.polls.lock().insert((fd, events)) as u64
+    }
+
+    fn poll(&self, poll_id: u64) -> Option<PollEvents> {
+        let poll_id = poll_id as usize;
+        let (fd, req_events) = *(self.polls.lock().get(poll_id)?);
+        if !req_events.contains(PollEvents::POLLIN) && !req_events.contains(PollEvents::POLLOUT) {
+            return None;
+        }
+        let pty_no = if fd as usize >= usize::MAX / 2 {
+            fd as usize - usize::MAX / 2
+        } else {
+           fd as usize
+        };
+        let slave = fd as usize >= usize::MAX / 2;
+        loop {
+            if !self.polls.lock().contains(poll_id) {
+                return None;
+            }
+            let curr_poll =if slave {
+                 self.ptys.read()[pty_no].curr_slave_poll()
+            } else {
+                 self.ptys.read()[pty_no].curr_master_poll()
+            };
+            if !((curr_poll & req_events).is_empty()) {
+                self.polls.lock().remove(poll_id);
+                return Some(curr_poll & req_events);
+            } else {
+                ipc::process_messages();
+            }
+        }
+    }
+
+    fn cancel_poll(&self, poll_id: u64) {
+        self.polls.lock().remove(poll_id as usize);
+    }
 }
 
 fn main() {
@@ -269,6 +321,7 @@ fn main() {
         ptys: Arc::new(RwLock::new(Slab::new())),
         slave_map: Arc::new(RwLock::new(HashMap::new())),
         devfs_client,
+        polls: Arc::new(Mutex::new(Slab::new())),
     };
     dev_driver_rpc::register_server(Box::new(serv.clone()));
     file_rpc::register_server(Box::new(serv.clone()));