Merge pull request #65 from SyedaAnshrahGillani/improve/xtask-server-and-signals

improve(xtask): add mdbook check, fix directory redirects, and use safe signal handling

Author: atulkhare4096

xtask/Cargo.toml

@@ -3,3 +3,6 @@ name = "xtask"
 version = "0.1.0"
 edition = "2021"
 publish = false
+
+[dependencies]
+ctrlc = "3.4"

xtask/src/main.rs

@@ -99,14 +99,32 @@ Commands:
 // ── build ────────────────────────────────────────────────────────────
 
 fn cmd_build() {
+    if !check_mdbook() {
+        eprintln!("Error: 'mdbook' not found in PATH. Please install it: https://rust-lang.github.io/mdbook/guide/installation.html");
+        std::process::exit(1);
+    }
     build_to("site");
 }
 
 fn cmd_deploy() {
+    if !check_mdbook() {
+        eprintln!("Error: 'mdbook' not found in PATH.");
+        std::process::exit(1);
+    }
     build_to("docs");
     println!("\nTo publish, commit docs/ and enable GitHub Pages → \"Deploy from a branch\" → /docs.");
 }
 
+fn check_mdbook() -> bool {
+    Command::new("mdbook")
+        .arg("--version")
+        .stdout(std::process::Stdio::null())
+        .stderr(std::process::Stdio::null())
+        .status()
+        .map(|s| s.success())
+        .unwrap_or(false)
+}
+
 fn build_to(dir_name: &str) {
     let root = project_root();
     let out = root.join(dir_name);
@@ -289,32 +307,68 @@ fn write_landing_page(site: &Path) {
     println!("  ✓ index.html");
 }
 
+enum ResolveResult {
+    File(PathBuf),
+    Redirect(String),
+    NotFound,
+}
+
 /// Resolve `request_target` (HTTP request path, e.g. `/foo/bar?x=1`) to a file under `site_canon`.
-/// Returns `None` for traversal attempts, missing files, or paths that escape `site_canon` (symlinks).
-fn resolve_site_file(site_canon: &Path, request_target: &str) -> Option<PathBuf> {
-    let path_only = request_target.split('?').next()?.split('#').next()?;
+/// Returns `ResolveResult::File` for success, `Redirect` if a trailing slash is needed for a directory,
+/// or `NotFound` for traversal attempts or missing files.
+///
+/// NOTE: This function preserves and hardens the multi-layer security from PR#18:
+/// 1. Percent-decoding via `percent_decode_path`.
+/// 2. Null byte rejection.
+/// 3. Traversal blocking (`..`).
+/// 4. Symlink escape prevention via canonicalization and prefix checking.
+fn resolve_site_file(site_canon: &Path, request_target: &str) -> ResolveResult {
+    let path_only = match request_target
+        .split('?')
+        .next()
+        .and_then(|s| s.split('#').next())
+    {
+        Some(p) => p,
+        None => return ResolveResult::NotFound,
+    };
+
+    // [Security] Handle percent-encoding and reject null bytes (from PR#18)
     let decoded = percent_decode_path(path_only);
     if decoded.as_bytes().contains(&0) {
-        return None;
+        return ResolveResult::NotFound;
     }
+
     let rel = decoded.trim_start_matches('/');
     let mut file_path = site_canon.to_path_buf();
     if !rel.is_empty() {
         for seg in rel.split('/').filter(|s| !s.is_empty()) {
+            // [Security] Block directory traversal (from PR#18)
             if seg == ".." {
-                return None;
+                return ResolveResult::NotFound;
             }
             file_path.push(seg);
         }
     }
+
     if file_path.is_dir() {
+        // If it refers to a directory but lacks a trailing slash, redirect so relative links work.
+        if !request_target.ends_with('/') && !request_target.is_empty() {
+            return ResolveResult::Redirect(format!("{path_only}/"));
+        }
         file_path.push("index.html");
     }
-    let real = fs::canonicalize(&file_path).ok()?;
-    if !real.starts_with(site_canon) {
-        return None;
+
+    // [Security] Canonicalize and verify we're still within site_canon (from PR#18)
+    let real = match fs::canonicalize(&file_path) {
+        Ok(r) => r,
+        Err(_) => return ResolveResult::NotFound,
+    };
+
+    if !real.starts_with(site_canon) || !real.is_file() {
+        return ResolveResult::NotFound;
     }
-    real.is_file().then_some(real)
+
+    ResolveResult::File(real)
 }
 
 fn hex_val(c: u8) -> Option<u8> {
@@ -360,79 +414,54 @@ fn cmd_serve() {
     println!("\nServing at http://{addr}  (Ctrl+C to stop)");
 
     for stream in listener.incoming() {
-        let Ok(stream) = stream else { continue };
-        handle_request(stream, &site_canon);
-    }
-}
-
-fn handle_request(mut stream: std::net::TcpStream, site_canon: &Path) {
-    let mut buf = [0u8; 4096];
-    let n = stream.read(&mut buf).unwrap_or(0);
-    let request = String::from_utf8_lossy(&buf[..n]);
-
-    let path = request
-        .lines()
-        .next()
-        .and_then(|line| line.split_whitespace().nth(1))
-        .unwrap_or("/");
-
-    if let Some(file_path) = resolve_site_file(site_canon, path) {
-        let body = fs::read(&file_path).unwrap_or_default();
-        let mime = guess_mime(&file_path);
-        let header = format!(
-            "HTTP/1.1 200 OK\r\nContent-Type: {mime}\r\nContent-Length: {}\r\n\r\n",
-            body.len()
-        );
-        let _ = stream.write_all(header.as_bytes());
-        let _ = stream.write_all(&body);
-    } else {
-        let body = b"404 Not Found";
-        let header = format!(
-            "HTTP/1.1 404 Not Found\r\nContent-Length: {}\r\n\r\n",
-            body.len()
-        );
-        let _ = stream.write_all(header.as_bytes());
-        let _ = stream.write_all(body);
+        let Ok(mut stream) = stream else { continue };
+        let mut buf = [0u8; 4096];
+        let n = stream.read(&mut buf).unwrap_or(0);
+        let request = String::from_utf8_lossy(&buf[..n]);
+
+        let path = request
+            .lines()
+            .next()
+            .and_then(|line| line.split_whitespace().nth(1))
+            .unwrap_or("/");
+
+        match resolve_site_file(&site_canon, path) {
+            ResolveResult::File(file_path) => {
+                let body = fs::read(&file_path).unwrap_or_default();
+                let mime = guess_mime(&file_path);
+                let header = format!(
+                    "HTTP/1.1 200 OK\r\nContent-Type: {mime}\r\nContent-Length: {}\r\n\r\n",
+                    body.len()
+                );
+                let _ = stream.write_all(header.as_bytes());
+                let _ = stream.write_all(&body);
+            }
+            ResolveResult::Redirect(new_path) => {
+                let header = format!(
+                    "HTTP/1.1 301 Moved Permanently\r\nLocation: {new_path}\r\nContent-Length: 0\r\n\r\n"
+                );
+                let _ = stream.write_all(header.as_bytes());
+            }
+            ResolveResult::NotFound => {
+                let body = b"404 Not Found";
+                let header = format!(
+                    "HTTP/1.1 404 Not Found\r\nContent-Length: {}\r\n\r\n",
+                    body.len()
+                );
+                let _ = stream.write_all(header.as_bytes());
+                let _ = stream.write_all(body);
+            }
+        }
     }
 }
 
 /// Install a Ctrl+C handler that exits cleanly (code 0) instead of
 /// letting the OS terminate with STATUS_CONTROL_C_EXIT.
 fn ctrlc_exit() {
-    unsafe {
-        libc_set_handler();
-    }
-}
-
-#[cfg(windows)]
-unsafe fn libc_set_handler() {
-    // SetConsoleCtrlHandler via the Windows API
-    extern "system" {
-        fn SetConsoleCtrlHandler(
-            handler: Option<unsafe extern "system" fn(u32) -> i32>,
-            add: i32,
-        ) -> i32;
-    }
-    unsafe extern "system" fn handler(_ctrl_type: u32) -> i32 {
+    ctrlc::set_handler(move || {
         std::process::exit(0);
-    }
-    unsafe {
-        SetConsoleCtrlHandler(Some(handler), 1);
-    }
-}
-
-#[cfg(not(windows))]
-unsafe fn libc_set_handler() {
-    // On Unix, register SIGINT via libc
-    extern "C" {
-        fn signal(sig: i32, handler: extern "C" fn(i32)) -> usize;
-    }
-    extern "C" fn handler(_sig: i32) {
-        std::process::exit(0);
-    }
-    unsafe {
-        signal(2 /* SIGINT */, handler);
-    }
+    })
+    .expect("Error setting Ctrl-C handler");
 }
 
 fn guess_mime(path: &Path) -> &'static str {
@@ -461,4 +490,4 @@ fn cmd_clean() {
             println!("Removed {dir_name}/");
         }
     }
-}
+}