From ff2a1ad26996cccdc9c6dcc6295702a34ebaa9d3 Mon Sep 17 00:00:00 2001
From: Toshit Chawda <r58playz@gmail.com>
Date: Mon, 25 Mar 2024 17:51:00 -0700
Subject: [PATCH] fix the prefix

---
 server/src/main.rs | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/server/src/main.rs b/server/src/main.rs
index 86c0661..5c90ce3 100644
--- a/server/src/main.rs
+++ b/server/src/main.rs
@@ -28,8 +28,8 @@ type HttpBody = http_body_util::Full<hyper::body::Bytes>;
 #[command(version = clap::crate_version!())]
 struct Cli {
     /// URL prefix the server should serve on
-    #[arg(long, default_value = "")]
-    prefix: String,
+    #[arg(long)]
+    prefix: Option<String>,
     /// Port the server should bind to
     #[arg(long, short, default_value = "4000")]
     port: String,
@@ -117,13 +117,18 @@ async fn main() -> Result<(), Error> {
 
     let socket = bind(&addr, opt.unix_socket).await?;
 
-    let prefix = if opt.prefix.starts_with('/') {
-        opt.prefix
+    let prefix = if let Some(prefix) = opt.prefix {
+        match (prefix.starts_with('/'), prefix.ends_with('/')) {
+            (true, true) => prefix,
+            (true, false) => prefix + "/",
+            (false, true) => "/".to_string() + &prefix,
+            (false, false) => "/".to_string() + &prefix + "/",
+        }
     } else {
-        "/".to_string() + &opt.prefix
+        "/".to_string()
     };
 
-    println!("listening on `{}`", addr);
+    println!("listening on `{}` with prefix `{}`", addr, prefix);
     while let Ok((stream, addr)) = socket.accept().await {
         let prefix = prefix.clone();
         tokio::spawn(async move {
@@ -155,7 +160,7 @@ async fn accept_http(
     {
         let (res, fut) = upgrade::upgrade(&mut req)?;
 
-        if uri == "/" {
+        if uri.is_empty() {
             tokio::spawn(async move { accept_ws(fut, addr.clone(), block_local).await });
         } else if let Some(uri) = uri.strip_prefix('/').map(|x| x.to_string()) {
             tokio::spawn(async move { accept_wsproxy(fut, uri, addr.clone(), block_local).await });