From 50b4e524b10d466a53733c3c05a79b41ae10eaee Mon Sep 17 00:00:00 2001
From: 4lDO2 <4lDO2@protonmail.com>
Date: Thu, 10 Apr 2025 21:00:26 +0200
Subject: [PATCH] Implement cancellation for netstack.

---
 netstack/src/smolnetd/scheme/socket.rs | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/netstack/src/smolnetd/scheme/socket.rs b/netstack/src/smolnetd/scheme/socket.rs
index 4ba2a27d6e..c1d1dc228f 100644
--- a/netstack/src/smolnetd/scheme/socket.rs
+++ b/netstack/src/smolnetd/scheme/socket.rs
@@ -156,6 +156,7 @@ where
 #[derive(Default, Clone)]
 struct WaitHandle {
     until: Option<TimeSpec>,
+    cancelling: bool,
     packet: SyscallPacket,
 }
 
@@ -288,7 +289,13 @@ where
                 }
             }
             if packet.a == KSMSG_CANCEL {
-                println!("smolnetd: todo: handle cancellation");
+                if let Some(idx) = self
+                    .wait_queue
+                    .iter()
+                    .position(|q| q.packet.id == packet.b as u64)
+                {
+                    self.wait_queue[idx].cancelling = true;
+                }
                 continue;
             }
             if let Some(a) = self.handle(&mut packet) {
@@ -299,6 +306,7 @@ where
                     Ok(timeout) => {
                         self.wait_queue.push(WaitHandle {
                             until: timeout,
+                            cancelling: false,
                             packet,
                         });
                     }
@@ -340,6 +348,12 @@ where
                 packet.a = a;
                 self.scheme_file.write_all(&packet)?;
             } else {
+                if self.wait_queue[i].cancelling {
+                    self.wait_queue.remove(i);
+                    packet.a = (-syscall::ECANCELED) as usize;
+                    self.scheme_file.write_all(&packet)?;
+                    continue;
+                }
                 match self.wait_queue[i].until {
                     Some(until)
                         if (until.tv_sec < cur_time.tv_sec
@@ -573,12 +587,13 @@ where
             socket.close_file(&scheme_file, &mut self.scheme_data)?;
         }
 
-        self.wait_queue.retain(
+        // incorrect, and kernel can't send close until all references are gone
+        /*self.wait_queue.retain(
             |&WaitHandle {
                  packet: SyscallPacket { a, .. },
                  ..
              }| a != fd,
-        );
+        );*/
 
         let remove = match self.ref_counts.entry(socket_handle) {
             Entry::Vacant(_) => {