optimize cstring reading

Reading cstrings byte by byte is very inefficient.

Author: Freax13

tee/kernel/src/memory/pagetable.rs

@@ -30,6 +30,7 @@ use constants::{
 use flush::{FlushGuard, GlobalFlushGuard};
 use log::trace;
 use static_page_tables::{StaticPageTable, StaticPd, StaticPdp, StaticPml4, flags};
+use usize_conversions::usize_from;
 use x86_64::{
     PhysAddr, VirtAddr,
     instructions::tlb::Pcid,
@@ -264,6 +265,44 @@ fn try_read_fast(src: VirtAddr, dest: NonNull<[u8]>) -> Result<(), ()> {
     if failed == 0 { Ok(()) } else { Err(()) }
 }
 
+/// Try to copy memory from `src` into `dest` until a null-byte is encountered.
+///
+/// Returns the number of read bytes.
+///
+/// This function is not unsafe. If the read fails for some reason `Err(())` is
+/// returned.
+#[inline(always)]
+fn try_read_cstring_fast(src: VirtAddr, dest: NonNull<[u8]>) -> Result<usize, usize> {
+    let remaining: u64;
+    let failed: u64;
+    unsafe {
+        asm!(
+            "test rcx, rcx",
+            "je 67f",
+            "66:",
+            "mov {scratch:l}, byte ptr [{src}]",
+            "mov byte ptr [{dest}], {scratch:l}",
+            "inc {src}",
+            "inc {dest}",
+            "test {scratch:l}, {scratch:l}",
+            "loopne 66b",
+            "67:",
+            ".pushsection .recoverable",
+            ".quad 66b",
+            ".quad 67b",
+            ".popsection",
+            src = inout(reg) src.as_u64() => _,
+            dest = inout(reg) dest.as_mut_ptr() => _,
+            inout("rcx") dest.len() => remaining,
+            scratch = out(reg) _,
+            inout("rdx") 0u64 => failed,
+        );
+    }
+
+    let read = dest.len() - usize_from(remaining);
+    if failed == 0 { Ok(read) } else { Err(read) }
+}
+
 /// Try to copy memory from `src` into `dest`.
 ///
 /// If the write fails for some reason `Err(())` is returned.
@@ -677,6 +716,23 @@ impl Pagetables {
         without_smap(|| try_read_fast(src, dest))
     }
 
+    /// Try to copy memory from `src` into `dest` until a null-byte is encountered.
+    ///
+    /// This function is not unsafe. If the read fails for some reason `Err(())` is
+    /// returned. If `src` isn't user memory `Err(())` is returned.
+    #[inline(always)]
+    pub fn try_read_cstring_user_fast(
+        &self,
+        src: VirtAddr,
+        dest: NonNull<[u8]>,
+    ) -> Result<usize, usize> {
+        check_user_address(src, dest.len()).map_err(|_| 0usize)?;
+
+        let _guard = self.activate();
+
+        without_smap(|| try_read_cstring_fast(src, dest))
+    }
+
     /// Try to copy memory from `src` into `dest`.
     ///
     /// If the write fails for some reason `Err(())` is returned. If `src` isn't

tee/kernel/src/user/process/memory.rs

@@ -297,6 +297,46 @@ impl VirtualMemory {
         Ok(ret)
     }
 
+    /// Read a string from userspace.
+    pub fn read_small_cstring(
+        &self,
+        pointer: Pointer<CString>,
+        max_length: usize,
+    ) -> Result<CString> {
+        let capacity = max_length + 2;
+        let mut buf = Vec::with_capacity(capacity);
+        let dest = unsafe { NonNull::new_unchecked(buf.as_mut_ptr()) };
+        let dest = NonNull::from_raw_parts(dest, capacity);
+
+        let addr = pointer.get();
+        check_user_address(addr, capacity)?;
+
+        let len = 'read: {
+            for _ in 0..Self::ACCESS_RETRIES {
+                match self.pagetables.try_read_cstring_user_fast(addr, dest) {
+                    Ok(len) => break 'read len,
+                    Err(read) => self.map_addrs(
+                        addr + u64::from_usize(read) + 1,
+                        1,
+                        PageTableFlags::empty(),
+                    )?,
+                }
+            }
+
+            self.pagetables
+                .try_read_cstring_user_fast(addr, dest)
+                .unwrap()
+        };
+
+        ensure!(len < capacity, NameTooLong);
+        unsafe {
+            buf.set_len(len);
+        }
+
+        let cstr = unsafe { CString::from_vec_with_nul_unchecked(buf) };
+        Ok(cstr)
+    }
+
     pub fn set_bytes(&self, addr: VirtAddr, count: usize, byte: u8) -> Result<()> {
         if count == 0 {
             return Ok(());

tee/kernel/src/user/process/syscall.rs

@@ -3672,7 +3672,7 @@ fn mount(
 ) -> SyscallResult {
     let _dev_name = virtual_memory.read(dev_name)?;
     let dir_name = virtual_memory.read(dir_name)?;
-    let r#type = virtual_memory.read_cstring(r#type, 0x10)?;
+    let r#type = virtual_memory.read_small_cstring(r#type, 0x10)?;
 
     let create_node = match r#type.as_bytes() {
         b"devtmpfs" => devtmpfs::new,

tee/kernel/src/user/process/syscall/args/pointee.rs

@@ -307,7 +307,7 @@ impl Pointee for CStr {
     fn display(f: &mut dyn fmt::Write, addr: VirtAddr, thread: &ThreadGuard) -> fmt::Result {
         let res = thread
             .virtual_memory()
-            .read_cstring(Pointer::from(addr), 1024);
+            .read_small_cstring(Pointer::from(addr), 1024);
         match res {
             Ok(value) => write!(f, "{value:?}"),
             Err(_) => write!(f, "{:#x} (invalid ptr)", addr.as_u64()),
@@ -343,7 +343,7 @@ impl AbiAgnosticPointee for Path {}
 
 impl ReadablePointee for Path {
     fn read(addr: VirtAddr, vm: &VirtualMemory, _abi: Abi) -> Result<(usize, Self)> {
-        let pathname = vm.read_cstring(Pointer::from(addr), PATH_MAX)?;
+        let pathname = vm.read_small_cstring(Pointer::from(addr), PATH_MAX)?;
         let len = pathname.to_bytes_with_nul().len();
         let value = Path::new(pathname.into_bytes())?;
         Ok((len, value))