Support multiple servers

README.md

@@ -12,12 +12,21 @@ Usage: lsp-ws-proxy [-l <listen>] [-s] [-r] [-v]
 Start WebSocket proxy for the LSP Server.
 Anything after the option delimiter is used to start the server.
 
+Multiple servers can be registered by separating each with an option delimiter,
+and using the query parameter `name` to specify the command name on connection.
+If no query parameter is present, the first one is started.
+
 Examples:
-  lsp-ws-proxy -- langserver
-  lsp-ws-proxy -- langserver --stdio
-  lsp-ws-proxy --listen 8888 -- langserver --stdio
-  lsp-ws-proxy --listen 0.0.0.0:8888 -- langserver --stdio
-  lsp-ws-proxy -l 8888 -- langserver --stdio
+  lsp-ws-proxy -- rust-analyzer
+  lsp-ws-proxy -- typescript-language-server --stdio
+  lsp-ws-proxy --listen 8888 -- rust-analyzer
+  lsp-ws-proxy --listen 0.0.0.0:8888 -- rust-analyzer
+  # Register multiple servers.
+  # Choose the server with query parameter `name` when connecting.
+  lsp-ws-proxy --listen 9999 --sync --remap \
+    -- typescript-language-server --stdio \
+    -- css-languageserver --stdio \
+    -- html-languageserver --stdio
 
 Options:
   -l, --listen      address or port to listen on (default: 0.0.0.0:9999)

src/api/proxy.rs

@@ -1,4 +1,4 @@
-use std::{process::Stdio, str::FromStr};
+use std::{convert::Infallible, process::Stdio, str::FromStr};
 
 use futures_util::{
     future::{select, Either},
@@ -14,18 +14,38 @@ use super::with_context;
 
 #[derive(Debug, Clone)]
 pub struct Context {
-    pub command: Vec<String>,
+    /// One or more commands to start a Language Server.
+    pub commands: Vec<Vec<String>>,
+    /// Write file on save.
     pub sync: bool,
+    /// Remap relative `source://` to absolute `file://`.
     pub remap: bool,
+    /// Project root.
     pub cwd: Url,
 }
 
+#[derive(Clone, Debug, serde::Deserialize)]
+struct Query {
+    /// The command name of the Language Server to start.
+    /// If not specified, the first one is started.
+    name: String,
+}
+
+fn with_optional_query() -> impl Filter<Extract = (Option<Query>,), Error = Infallible> + Clone {
+    warp::query::<Query>()
+        .map(Some)
+        .or_else(|_| async { Ok::<(Option<Query>,), Infallible>((None,)) })
+}
+
 /// Handler for WebSocket connection.
 pub fn handler(ctx: Context) -> impl Filter<Extract = impl Reply, Error = Rejection> + Clone {
     warp::path::end()
         .and(warp::ws())
         .and(with_context(ctx))
-        .map(|ws: warp::ws::Ws, ctx| ws.on_upgrade(move |socket| on_upgrade(socket, ctx)))
+        .and(with_optional_query())
+        .map(|ws: warp::ws::Ws, ctx, query| {
+            ws.on_upgrade(move |socket| on_upgrade(socket, ctx, query))
+        })
 }
 
 #[tracing::instrument(level = "debug", err, skip(msg))]
@@ -47,27 +67,42 @@ async fn maybe_write_text_document(msg: &lsp::Message) -> Result<(), std::io::Er
     Ok(())
 }
 
-async fn on_upgrade(socket: warp::ws::WebSocket, ctx: Context) {
+async fn on_upgrade(socket: warp::ws::WebSocket, ctx: Context, query: Option<Query>) {
     tracing::info!("connected");
-    if let Err(err) = connected(socket, ctx).await {
+    if let Err(err) = connected(socket, ctx, query).await {
         tracing::error!("connection error: {}", err);
     }
     tracing::info!("disconnected");
 }
 
-#[tracing::instrument(level = "debug", skip(ws, ctx), fields(command = ?ctx.command[0], remap = %ctx.remap, sync = %ctx.sync))]
+#[tracing::instrument(level = "debug", skip(ws, ctx), fields(remap = %ctx.remap, sync = %ctx.sync))]
 async fn connected(
     ws: warp::ws::WebSocket,
     ctx: Context,
+    query: Option<Query>,
 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
-    tracing::info!("starting {} in {}", ctx.command[0], ctx.cwd);
-    let mut server = Command::new(&ctx.command[0])
-        .args(&ctx.command[1..])
+    let command = if let Some(query) = query {
+        if let Some(command) = ctx.commands.iter().find(|v| v[0] == query.name) {
+            command
+        } else {
+            // TODO Validate this earlier and reject, or close immediately.
+            tracing::warn!(
+                "Unknown Language Server '{}', falling back to the default",
+                query.name
+            );
+            &ctx.commands[0]
+        }
+    } else {
+        &ctx.commands[0]
+    };
+    tracing::info!("starting {} in {}", command[0], ctx.cwd);
+    let mut server = Command::new(&command[0])
+        .args(&command[1..])
         .stdin(Stdio::piped())
         .stdout(Stdio::piped())
         .kill_on_drop(true)
         .spawn()?;
-    tracing::debug!("running {}", ctx.command[0]);
+    tracing::debug!("running {}", command[0]);
 
     let mut server_send = lsp::framed::writer(server.stdin.take().unwrap());
     let mut server_recv = lsp::framed::reader(server.stdout.take().unwrap());

src/main.rs

@@ -14,12 +14,21 @@ mod lsp;
 Start WebSocket proxy for the LSP Server.
 Anything after the option delimiter is used to start the server.
 
+Multiple servers can be registered by separating each with an option delimiter,
+and using the query parameter `name` to specify the command name on connection.
+If no query parameter is present, the first one is started.
+
 Examples:
-  lsp-ws-proxy -- langserver
-  lsp-ws-proxy -- langserver --stdio
-  lsp-ws-proxy --listen 8888 -- langserver --stdio
-  lsp-ws-proxy --listen 0.0.0.0:8888 -- langserver --stdio
-  lsp-ws-proxy -l 8888 -- langserver --stdio
+  lsp-ws-proxy -- rust-analyzer
+  lsp-ws-proxy -- typescript-language-server --stdio
+  lsp-ws-proxy --listen 8888 -- rust-analyzer
+  lsp-ws-proxy --listen 0.0.0.0:8888 -- rust-analyzer
+  # Register multiple servers.
+  # Choose the server with query parameter `name` when connecting.
+  lsp-ws-proxy --listen 9999 --sync --remap \
+    -- typescript-language-server --stdio \
+    -- css-languageserver --stdio \
+    -- html-languageserver --stdio
 */
 struct Options {
     /// address or port to listen on (default: 0.0.0.0:9999)
@@ -47,7 +56,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         .with_env_filter(std::env::var("RUST_LOG").unwrap_or_else(|_| "info".to_owned()))
         .init();
 
-    let (opts, command) = get_opts_and_command();
+    let (opts, commands) = get_opts_and_commands();
 
     let cwd = std::env::current_dir()?;
     // TODO Move these to `api` module.
@@ -58,7 +67,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     // TODO Limit concurrent connection. Can get messy when `sync` is used.
     // TODO? Keep track of added files and remove them on disconnect?
     let proxy = api::proxy::handler(api::proxy::Context {
-        command,
+        commands,
         sync: opts.sync,
         remap: opts.remap,
         cwd: Url::from_directory_path(&cwd).expect("valid url from current dir"),
@@ -82,9 +91,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     Ok(())
 }
 
-fn get_opts_and_command() -> (Options, Vec<String>) {
-    let strings: Vec<String> = std::env::args().collect();
-    let splitted: Vec<&[String]> = strings.splitn(2, |s| *s == "--").collect();
+fn get_opts_and_commands() -> (Options, Vec<Vec<String>>) {
+    let args: Vec<String> = std::env::args().collect();
+    let splitted: Vec<Vec<String>> = args.split(|s| *s == "--").map(|s| s.to_vec()).collect();
     let strs: Vec<&str> = splitted[0].iter().map(|s| s.as_str()).collect();
 
     // Parse options or show help and exit.
@@ -102,11 +111,12 @@ fn get_opts_and_command() -> (Options, Vec<String>) {
         std::process::exit(0);
     }
 
-    if splitted.len() != 2 {
+    if splitted.len() < 2 {
         panic!("Command to start the server is required. See --help for examples.");
     }
 
-    (opts, splitted[1].to_vec())
+    let commands = splitted[1..].iter().map(|s| s.to_owned()).collect();
+    (opts, commands)
 }
 
 fn parse_listen(value: &str) -> Result<String, String> {