From 0029070acbbdb1c51f9ed5defda082a56c3a685b Mon Sep 17 00:00:00 2001
From: Jiajie Chen <noc@jiegec.ac.cn>
Date: Fri, 15 Mar 2019 08:55:58 +0800
Subject: [PATCH] Fix rust/sh and add rust/tls test

---
 rust/src/bin/sh.rs  | 15 +++++++++----
 rust/src/bin/tls.rs | 52 +++++++++++++++++++++++++++++++++++++++++++++
 rust/src/syscall.rs | 18 +++++++++++++++-
 3 files changed, 80 insertions(+), 5 deletions(-)
 create mode 100644 rust/src/bin/tls.rs

diff --git a/rust/src/bin/sh.rs b/rust/src/bin/sh.rs
index 90cc91d..d9482e5 100644
--- a/rust/src/bin/sh.rs
+++ b/rust/src/bin/sh.rs
@@ -22,9 +22,12 @@ pub fn main() -> i32 {
         print!(">> ");
         let cmd = get_line(&mut history);
         // split cmd, make argc & argv
-        let cmd = cmd.replace(' ', "\0");
-        let mut ptrs: Vec<usize> = cmd.split('\0')
-            .filter(|s| !s.is_empty()).map(|s| s.as_ptr() as usize).collect();
+        let cmd = cmd.replace(' ', "\0") + "\0";
+        let mut ptrs: Vec<usize> = cmd
+            .split('\0')
+            .filter(|s| !s.is_empty())
+            .map(|s| s.as_ptr() as usize)
+            .collect();
         if ptrs.is_empty() {
             continue;
         }
@@ -33,7 +36,11 @@ pub fn main() -> i32 {
         let pid = sys_fork();
         assert!(pid >= 0);
         if pid == 0 {
-            return sys_exec(ptrs[0] as *const u8, ptrs.as_ptr() as *const *const u8, ptr::null());
+            return sys_exec(
+                ptrs[0] as *const u8,
+                ptrs.as_ptr() as *const *const u8,
+                ptr::null(),
+            );
         } else {
             let mut code: i32 = 0;
             sys_wait(pid as usize, &mut code);
diff --git a/rust/src/bin/tls.rs b/rust/src/bin/tls.rs
new file mode 100644
index 0000000..9646e87
--- /dev/null
+++ b/rust/src/bin/tls.rs
@@ -0,0 +1,52 @@
+#![no_std]
+#![no_main]
+#![feature(asm)]
+
+#[macro_use]
+extern crate rcore_user;
+
+use rcore_user::syscall::{sys_arch_prctl, sys_fork, sys_getpid, sys_sleep};
+
+fn set_tls(tls: usize, pid: usize) {
+    #[cfg(any(target_arch = "riscv32", target_arch = "riscv64"))]
+    unsafe {
+        asm!("mv tp, $0" : : "r"(tls));
+    }
+    #[cfg(target_arch = "x86_64")]
+    unsafe {
+        static mut DATA: [usize; 1024] = [0; 1024];
+        DATA[pid] = tls;
+        // set fs base
+        sys_arch_prctl(0x1002, &DATA[pid] as *const usize as usize);
+    }
+}
+
+fn get_tls() -> usize {
+    let mut tls: usize = 0;
+    #[cfg(any(target_arch = "riscv32", target_arch = "riscv64"))]
+    unsafe {
+        asm!("mv $0, tp" : "=r"(tls) :);
+    }
+    #[cfg(target_arch = "x86_64")]
+    unsafe {
+        asm!("mov %fs:0, $0" : "=r"(tls) :);
+    }
+    tls
+}
+
+// IMPORTANT: Must define main() like this
+#[no_mangle]
+pub fn main() {
+    println!("I am going to set TLS specific registers");
+    println!("And see if it changes between context switch");
+    sys_fork();
+    sys_fork();
+    let pid = sys_getpid();
+    for i in 0..10 {
+        let magic = 0xcafebabe;
+        set_tls(magic + i, pid as usize);
+        sys_sleep(1);
+        assert_eq!(get_tls(), magic + i);
+        println!("pid {} time {}: working", pid, i);
+    }
+}
diff --git a/rust/src/syscall.rs b/rust/src/syscall.rs
index c600014..d340aac 100644
--- a/rust/src/syscall.rs
+++ b/rust/src/syscall.rs
@@ -93,8 +93,19 @@ pub fn sys_getpid() -> i32 {
     sys_call(SyscallId::GetPid, 0, 0, 0, 0, 0, 0)
 }
 
+#[repr(C)]
+#[derive(Debug, Copy, Clone)]
+pub struct TimeSpec {
+    sec: u64,
+    nsec: u64,
+}
+
 pub fn sys_sleep(time: usize) -> i32 {
-    sys_call(SyscallId::Sleep, time, 0, 0, 0, 0, 0)
+    let ts = TimeSpec {
+        sec: time as u64,
+        nsec: 0
+    };
+    sys_call(SyscallId::Sleep, &ts as *const TimeSpec as usize, 0, 0, 0, 0, 0)
 }
 
 pub fn sys_get_time() -> i32 {
@@ -105,6 +116,10 @@ pub fn sys_set_priority(priority: usize) -> i32 {
     sys_call(SyscallId::SetPriority, priority, 0, 0, 0, 0, 0)
 }
 
+pub fn sys_arch_prctl(code: i32, addr: usize) -> i32 {
+    sys_call(SyscallId::ArchPrctl, code as usize, addr, 0, 0, 0, 0)
+}
+
 #[allow(dead_code)]
 enum SyscallId {
     Read = 0,
@@ -130,4 +145,5 @@ enum SyscallId {
     GetCwd = 79,
     GetTime = 96,
     SetPriority = 141,
+    ArchPrctl = 158,
 }