From 113583cfca1f55e936a1e869b019d37aa09a56ce Mon Sep 17 00:00:00 2001
From: 4lDO2 <4lDO2@protonmail.com>
Date: Tue, 4 Jul 2023 18:34:50 +0200
Subject: [PATCH] Enforce page alignment in shm:, better fmap.

---
 Cargo.lock  |  8 ++---
 Cargo.toml  |  2 +-
 src/main.rs |  4 ++-
 src/shm.rs  | 85 +++++++++++++++++++++++++++++++++++++++++------------
 4 files changed, 74 insertions(+), 25 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index c41c13af03..9bffb936bd 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -18,15 +18,13 @@ dependencies = [
 
 [[package]]
 name = "libc"
-version = "0.2.139"
+version = "0.2.147"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "201de327520df007757c1f0adce6e827fe8562fbc28bfd9c15571c66ca1f5f79"
+checksum = "b4668fb0ea861c1df094127ac5f1da3409a82116a4ba74fca2e58ef927159bb3"
 
 [[package]]
 name = "redox_syscall"
-version = "0.2.16"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fb5a58c1855b4b6819d59012155603f0b22ad30cad752600aadfcb695265519a"
+version = "0.3.5"
 dependencies = [
  "bitflags",
 ]
diff --git a/Cargo.toml b/Cargo.toml
index e3a780c384..846a17fd2b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -4,4 +4,4 @@ name = "ipcd"
 version = "0.1.0"
 [dependencies]
 libc = "0.2"
-redox_syscall = "0.2"
+redox_syscall = { path = "../../kernel/source/syscall" }
diff --git a/src/main.rs b/src/main.rs
index ba556599c1..0c45d7d102 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -112,7 +112,9 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
                         match shm.socket.read(&mut packet) {
                             Ok(0) => break true,
                             Ok(_) => {
-                                shm.handle(&mut packet);
+                                unsafe {
+                                    shm.do_handle(&mut packet);
+                                }
                                 shm.socket.write(&packet)?;
                             },
                             Err(err) => if err.kind() == io::ErrorKind::WouldBlock {
diff --git a/src/shm.rs b/src/shm.rs
index 92a1bdb684..f1855e1ee9 100644
--- a/src/shm.rs
+++ b/src/shm.rs
@@ -6,11 +6,11 @@ use std::{
     os::unix::fs::OpenOptionsExt,
     rc::Rc,
 };
-use syscall::{error::*, flag::O_NONBLOCK, Error, Map, SchemeMut, Result};
+use syscall::{error::*, flag::O_NONBLOCK, Error, Map, SchemeMut, Result, Packet, KSMSG_MMAP_PREP, MapFlags, KSMSG_MMAP, PAGE_SIZE, MAP_PRIVATE};
 
 #[derive(Default)]
 pub struct ShmHandle {
-    buffer: Option<Box<[u8]>>,
+    buffer: Option<MmapGuard>,
     refs: usize
 }
 pub struct ShmScheme {
@@ -33,6 +33,41 @@ impl ShmScheme {
                 .open(":shm")?
         })
     }
+    pub unsafe fn do_handle(&mut self, packet: &mut Packet) {
+        match packet.a {
+            KSMSG_MMAP_PREP | KSMSG_MMAP => {
+                let req_file = packet.b;
+                let req_flags = MapFlags::from_bits_truncate(packet.c);
+                let req_page_count = packet.d;
+                let req_offset = u64::from(packet.uid) | (u64::from(packet.gid) << 32);
+
+                let res = self.ksmsg_mmap(req_file, req_flags, req_offset, req_page_count);
+
+                *packet = Packet {
+                    id: packet.id,
+                    a: syscall::Error::mux(res),
+                    ..Packet::default()
+                };
+            }
+            _ => self.handle(packet),
+        }
+    }
+    pub fn ksmsg_mmap(&mut self, id: usize, _flags: MapFlags, offset: u64, page_count: usize) -> Result<usize> {
+        let path = self.handles.get(&id).ok_or(Error::new(EBADF))?;
+        let total_size = offset as usize + page_count * PAGE_SIZE;
+        match self.maps.get_mut(path).expect("handle pointing to nothing").buffer {
+            Some(ref mut buf) => {
+                if total_size > buf.len() {
+                    return Err(Error::new(ERANGE));
+                }
+                Ok(buf.as_ptr() + offset as usize)
+            },
+            ref mut buf @ None => {
+                *buf = Some(MmapGuard::alloc(page_count)?);
+                Ok(buf.as_mut().unwrap().as_ptr() + offset as usize)
+            }
+        }
+    }
 }
 
 impl SchemeMut for ShmScheme {
@@ -46,22 +81,6 @@ impl SchemeMut for ShmScheme {
         self.next_id += 1;
         Ok(id)
     }
-    fn fmap(&mut self, id: usize, map: &Map) -> Result<usize> {
-        let path = self.handles.get(&id).ok_or(Error::new(EBADF))?;
-        let total_size = map.offset + map.size;
-        match self.maps.get_mut(path).expect("handle pointing to nothing").buffer {
-            Some(ref mut buf) => {
-                if total_size != buf.len() {
-                    return Err(Error::new(ERANGE));
-                }
-                Ok(buf[map.offset..].as_mut_ptr() as usize)
-            },
-            ref mut buf @ None => {
-                *buf = Some(vec![0; total_size].into_boxed_slice());
-                Ok(buf.as_mut().unwrap()[map.offset..].as_mut_ptr() as usize)
-            }
-        }
-    }
     fn fpath(&mut self, id: usize, buf: &mut [u8]) -> Result<usize> {
         // Write scheme name
         const PREFIX: &[u8] = b"shm:";
@@ -92,3 +111,33 @@ impl SchemeMut for ShmScheme {
         Ok(0)
     }
 }
+
+pub struct MmapGuard {
+    base: usize,
+    size: usize,
+}
+impl MmapGuard {
+    pub fn alloc(page_count: usize) -> Result<Self> {
+        let size = page_count * PAGE_SIZE;
+        let base = unsafe { syscall::fmap(!0, &Map { offset: 0, size, flags: MAP_PRIVATE, address: 0 }) }?;
+
+        Ok(Self {
+            base,
+            size,
+        })
+    }
+    pub fn len(&self) -> usize {
+        self.size
+    }
+    pub fn as_ptr(&self) -> usize {
+        self.base
+    }
+}
+impl Drop for MmapGuard {
+    fn drop(&mut self) {
+        if self.size == 0 {
+            return;
+        }
+        let _ = unsafe { syscall::funmap(self.base, self.size) };
+    }
+}