fix message loss with requests with multiple responses

rustfmt.toml

@@ -0,0 +1,10 @@
+max_width = 120
+hard_tabs = false
+tab_spaces = 4
+edition = "2021"
+reorder_imports = true
+reorder_modules = true
+format_strings = true
+normalize_comments = true
+wrap_comments = true
+comment_width = 80

src/commands/args.rs

@@ -80,11 +80,7 @@ mod tests {
     fn parses_double_quoted_argument() {
         assert_eq!(
             split_command_line(r#"/cmd arg1 "arg2 value""#),
-            vec![
-                "/cmd".to_string(),
-                "arg1".to_string(),
-                "arg2 value".to_string()
-            ]
+            vec!["/cmd".to_string(), "arg1".to_string(), "arg2 value".to_string()]
         );
     }
 

src/commands/mod.rs

@@ -41,10 +41,7 @@ pub async fn global_usage(ctx: Arc<Mutex<BotContext>>) -> String {
 
     usage += "\nCommands from plugins:\n";
     for cmd in ctx_lock.plugin_commands.values() {
-        usage += &format!(
-            "{} (from {}) - {}\n",
-            &cmd.usage, &cmd.plugin_id, &cmd.description
-        );
+        usage += &format!("{} (from {}) - {}\n", &cmd.usage, &cmd.plugin_id, &cmd.description);
     }
 
     usage += "\nDeltaChat Remote Control Bot by slavasil, 2026\n";
@@ -65,12 +62,7 @@ pub async fn auth_command(
     let dchat_ctx_lock = dchat_ctx.lock().await;
 
     if !AUTH_REQUIRED {
-        chat::send_text_msg(
-            &dchat_ctx_lock,
-            chat_id,
-            "Authentication is disabled".to_owned(),
-        )
-        .await?;
+        chat::send_text_msg(&dchat_ctx_lock, chat_id, "Authentication is disabled".to_owned()).await?;
         return Ok(());
     }
 
@@ -80,23 +72,13 @@ pub async fn auth_command(
     }
 
     let Some(password) = args.first() else {
-        chat::send_text_msg(
-            &dchat_ctx_lock,
-            chat_id,
-            "Usage: /auth <password>".to_owned(),
-        )
-        .await?;
+        chat::send_text_msg(&dchat_ctx_lock, chat_id, "Usage: /auth <password>".to_owned()).await?;
         return Ok(());
     };
 
     if *password == ctx_lock.config.auth.password {
         ctx_lock.authed_contacts.insert(contact_id);
-        chat::send_text_msg(
-            &dchat_ctx_lock,
-            chat_id,
-            "Authentication successful!".to_owned(),
-        )
-        .await?;
+        chat::send_text_msg(&dchat_ctx_lock, chat_id, "Authentication successful!".to_owned()).await?;
     } else {
         chat::send_text_msg(&dchat_ctx_lock, chat_id, "Incorrect password".to_owned()).await?;
     }
@@ -201,12 +183,7 @@ async fn send_magic_packet(mac: MacAddress) -> AnyhowResult<()> {
 
 fn ssh_unlock_disk_usage(config: &BotConfig) -> String {
     if !config.machines.is_empty() {
-        let list = config
-            .machines
-            .keys()
-            .cloned()
-            .collect::<Vec<_>>()
-            .join("|");
+        let list = config.machines.keys().cloned().collect::<Vec<_>>().join("|");
         if config.default_machine().is_some() {
             format!("Usage: /ssh-unlock-disk [{list}] <password>")
         } else {
@@ -235,12 +212,7 @@ pub async fn ssh_unlock_disk_command(
 
     if args.is_empty() || args.len() > 2 {
         log::warn!("wrong number of args");
-        chat::send_text_msg(
-            &dchat_ctx_lock,
-            chat_id,
-            ssh_unlock_disk_usage(&ctx_lock.config),
-        )
-        .await?;
+        chat::send_text_msg(&dchat_ctx_lock, chat_id, ssh_unlock_disk_usage(&ctx_lock.config)).await?;
         return Ok(());
     }
 
@@ -255,12 +227,7 @@ pub async fn ssh_unlock_disk_command(
             None
         }
     }) else {
-        chat::send_text_msg(
-            &dchat_ctx_lock,
-            chat_id,
-            ssh_unlock_disk_usage(&ctx_lock.config),
-        )
-        .await?;
+        chat::send_text_msg(&dchat_ctx_lock, chat_id, ssh_unlock_disk_usage(&ctx_lock.config)).await?;
         return Ok(());
     };
 
@@ -284,8 +251,7 @@ pub async fn ssh_unlock_disk_command(
 
     let password = args[args.len() - 1].to_owned();
     let host = machine.static_ip.to_string();
-    let ssh_key =
-        russh::keys::PrivateKey::from_openssh(ssh_key).context("Invalid remote unlock SSH key")?;
+    let ssh_key = russh::keys::PrivateKey::from_openssh(ssh_key).context("Invalid remote unlock SSH key")?;
 
     drop(ctx_lock);
     drop(dchat_ctx_lock);
@@ -311,10 +277,7 @@ pub async fn ssh_unlock_disk_command(
             SSH_TIMEOUT,
             connection.authenticate_publickey(
                 "root",
-                PrivateKeyWithHashAlg::new(
-                    Arc::new(ssh_key),
-                    connection.best_supported_rsa_hash().await?.flatten(),
-                ),
+                PrivateKeyWithHashAlg::new(Arc::new(ssh_key), connection.best_supported_rsa_hash().await?.flatten()),
             ),
         )
         .await
@@ -350,10 +313,7 @@ pub async fn ssh_unlock_disk_command(
             .context("Cannot send password via SSH")?;
 
         loop {
-            let Some(msg) = timeout(SSH_TIMEOUT, channel.wait())
-                .await
-                .context("Timed out")?
-            else {
+            let Some(msg) = timeout(SSH_TIMEOUT, channel.wait()).await.context("Timed out")? else {
                 break;
             };
             log::info!("Received message from SSH server: {:?}", msg);
@@ -372,12 +332,7 @@ pub async fn ssh_unlock_disk_command(
 
 fn ssh_exec_usage(config: &BotConfig) -> String {
     if !config.machines.is_empty() {
-        let list = config
-            .machines
-            .keys()
-            .cloned()
-            .collect::<Vec<_>>()
-            .join("|");
+        let list = config.machines.keys().cloned().collect::<Vec<_>>().join("|");
         format!("Usage: /ssh-exec <username>@({list}|<hostname>|<IP address>) <command...>")
     } else {
         "Command /ssh-exec unavailable: no machines configured".to_string()
@@ -499,24 +454,16 @@ pub async fn ssh_exec_command(
                     russh::ChannelMsg::Data { data } => {
                         let dchat_ctx_lock = dchat_ctx.lock().await;
                         report_text += &String::from_utf8_lossy(&data);
-                        chat::send_edit_request(
-                            &dchat_ctx_lock,
-                            report_message_id,
-                            report_text.clone(),
-                        )
-                        .await
-                        .context("Cannot edit message")?;
+                        chat::send_edit_request(&dchat_ctx_lock, report_message_id, report_text.clone())
+                            .await
+                            .context("Cannot edit message")?;
                     }
                     russh::ChannelMsg::ExtendedData { data, ext: 1 } => {
                         let dchat_ctx_lock = dchat_ctx.lock().await;
                         report_text += &String::from_utf8_lossy(&data);
-                        chat::send_edit_request(
-                            &dchat_ctx_lock,
-                            report_message_id,
-                            report_text.clone(),
-                        )
-                        .await
-                        .context("Cannot edit message")?;
+                        chat::send_edit_request(&dchat_ctx_lock, report_message_id, report_text.clone())
+                            .await
+                            .context("Cannot edit message")?;
                     }
                     _ => {
                         log::info!("Message from SSH channel: {msg:?}");
@@ -531,13 +478,9 @@ pub async fn ssh_exec_command(
 
                     let dchat_ctx_lock = dchat_ctx.lock().await;
                     report_text += "\n\nRead timed out";
-                    chat::send_edit_request(
-                        &dchat_ctx_lock,
-                        report_message_id,
-                        report_text.clone(),
-                    )
-                    .await
-                    .context("Cannot edit message")?;
+                    chat::send_edit_request(&dchat_ctx_lock, report_message_id, report_text.clone())
+                        .await
+                        .context("Cannot edit message")?;
 
                     break;
                 }

src/config/util.rs

@@ -7,78 +7,107 @@ pub(crate) fn yaml_to_json(yaml: &Value) -> String {
 }
 
 fn write_yaml_to_json(yaml: &Value, out: &mut String) {
+    eprintln!("write_yaml_to_json {:?}", yaml);
     match yaml {
         Value::Null => {
+            eprintln!("write null");
             out.push_str("null");
         }
         Value::Bool(b) => {
+            eprintln!("write true/false");
             out.push_str(if *b { "true" } else { "false" });
         }
         Value::Number(n) => {
+            eprintln!("write {n}");
             out.push_str(&n.to_string());
         }
         Value::String(s) => {
+            eprintln!("write \"");
             out.push('"');
             for c in s.encode_utf16() {
                 match c {
                     0x005C | 0x0022 => {
+                        eprintln!("write \\{}", char::from_u32(c as u32).unwrap());
                         out.push('\\');
                         out.push(char::from_u32(c as u32).unwrap());
                     }
                     0x0008 => {
+                        eprintln!("write \\b");
                         out.push('\\');
                         out.push('b');
                     }
                     0x0009 => {
+                        eprintln!("write \\t");
                         out.push('\\');
                         out.push('t');
                     }
                     0x000a => {
+                        eprintln!("write \\n");
                         out.push('\\');
                         out.push('n');
                     }
                     0x000C => {
+                        eprintln!("write \\f");
                         out.push('\\');
                         out.push('f');
                     }
                     0x0000..=0x001F | 0x0080..=0xFFFF => {
+                        eprintln!("write \\u{c:04x}");
                         out.push_str(&format!("\\u{c:04x}"));
                     }
                     _ => {
+                        eprintln!("write {}", char::from_u32(c as u32).unwrap());
                         out.push(char::from_u32(c as u32).unwrap());
                     }
                 }
             }
-            out.push_str(s); // TODO escape
+            eprintln!("write \"");
             out.push('"');
         }
         Value::Tagged(_) => {}
         Value::Sequence(list) => {
+            eprintln!("write [");
             out.push('[');
             if !list.is_empty() {
                 write_yaml_to_json(&list[0], out);
                 for i in 1..list.len() {
+                    eprintln!("write ,");
                     out.push(',');
                     write_yaml_to_json(&list[i], out);
                 }
             }
+            eprintln!("write ]");
             out.push(']');
         }
         Value::Mapping(map) => {
+            eprintln!("write {{");
             out.push('{');
             let mut iter = map.iter();
             if let Some(kv) = iter.next() {
-                write_yaml_to_json(&Value::String(yaml_to_json(kv.0)), out);
+                let key = match kv.0 {
+                    Value::String(_) => kv.0,
+                    _ => &Value::String(yaml_to_json(kv.0)),
+                };
+                write_yaml_to_json(key, out);
+                eprintln!("write :");
                 out.push(':');
                 write_yaml_to_json(kv.1, out);
             }
             for kv in iter {
+                eprintln!("write ,");
                 out.push(',');
-                write_yaml_to_json(&Value::String(yaml_to_json(kv.0)), out);
+                let key = match kv.0 {
+                    Value::String(_) => kv.0,
+                    _ => &Value::String(yaml_to_json(kv.0)),
+                };
+                write_yaml_to_json(key, out);
+                eprintln!("write :");
                 out.push(':');
                 write_yaml_to_json(kv.1, out);
             }
+            eprintln!("write }}");
             out.push('}');
         }
     }
+    eprintln!("written {:?}", yaml);
 }

src/main.rs

@@ -7,10 +7,7 @@ mod ssh;
 mod proto {
     pub(crate) mod deltachat_remotecontrol_bot {
         pub(crate) mod plugin {
-            include!(concat!(
-                env!("OUT_DIR"),
-                "/deltachat_remotecontrol_bot.plugin.rs"
-            ));
+            include!(concat!(env!("OUT_DIR"), "/deltachat_remotecontrol_bot.plugin.rs"));
         }
     }
     pub(crate) use deltachat_remotecontrol_bot::plugin::*;
@@ -82,14 +79,9 @@ async fn run_bot(bot_context: Arc<Mutex<BotContext>>) -> AnyhowResult<()> {
 
     let dchat_db_filename = dchat_db_dir.join("deltachat.sqlite");
 
-    let dchat_ctx = Context::new(
-        &dchat_db_filename,
-        0,
-        Events::default(),
-        StockStrings::default(),
-    )
-    .await
-    .context("Failed to open Delta Chat client DB")?;
+    let dchat_ctx = Context::new(&dchat_db_filename, 0, Events::default(), StockStrings::default())
+        .await
+        .context("Failed to open Delta Chat client DB")?;
 
     let ctx_lock = bot_context.lock().await;
 
@@ -114,9 +106,7 @@ async fn run_bot(bot_context: Arc<Mutex<BotContext>>) -> AnyhowResult<()> {
                 avatar_path = data_path().join(&avatar_path);
             }
             if let Some(avatar_str) = avatar_path.to_str() {
-                dchat_ctx
-                    .set_config(Config::Selfavatar, Some(avatar_str))
-                    .await?;
+                dchat_ctx.set_config(Config::Selfavatar, Some(avatar_str)).await?;
             }
         }
     }
@@ -156,12 +146,7 @@ async fn send_command_not_found(
     chat_id: ChatId,
 ) -> AnyhowResult<()> {
     let dchat_ctx_lock = dchat_ctx.lock().await;
-    chat::send_text_msg(
-        &dchat_ctx_lock,
-        chat_id,
-        commands::global_usage(bot_ctx).await,
-    )
-    .await?;
+    chat::send_text_msg(&dchat_ctx_lock, chat_id, commands::global_usage(bot_ctx).await).await?;
     Ok(())
 }
 
@@ -208,9 +193,7 @@ async fn handle_message(
             };
 
             let Some(plugin) = ctx_lock.plugins.get(&cmd.plugin_id) else {
-                bail!(
-                    "Inconsistency in plugin_cmd_aliases hashmap: command references a plugin that is not loaded"
-                );
+                bail!("Inconsistency in plugin_cmd_aliases hashmap: command references a plugin that is not loaded");
             };
             let (plugin_conn, plugin_name) = {
                 let lock = plugin.lock().await;
@@ -218,11 +201,7 @@ async fn handle_message(
                     bail!("Plugin disconnected");
                 };
                 let plugin_name = lock.plugin_id.clone();
-                log::info!(
-                    "Delegating command /{} to plugin {}",
-                    &cmd.name,
-                    &plugin_name
-                );
+                log::info!("Delegating command /{} to plugin {}", &cmd.name, &plugin_name);
                 (conn, plugin_name)
             };
             let contact_id = msg.get_from_id();
@@ -239,16 +218,13 @@ async fn handle_message(
             tokio::spawn(async move {
                 log::debug!("Started plugin command task");
                 match plugin_conn
-                    .execute_command(
-                        cmd_name.clone(),
-                        issuer_id,
-                        commands::args::split_command_line(&text),
-                    )
+                    .execute_command(cmd_name.clone(), issuer_id, commands::args::split_command_line(&text))
                     .await
                 {
                     Ok(mut replies) => {
                         let mut last_message_id = None;
                         while let Some(reply) = replies.recv().await {
+                            log::trace!("reply recv'ed: Some({:?})", reply);
                             match reply {
                                 Ok(reply) => {
                                     log::debug!(
@@ -259,21 +235,15 @@ async fn handle_message(
                                     );
                                     if reply.edit {
                                         if let Some(msg_id) = last_message_id {
-                                            if let Some(proto::command_reply::Reply::Text(text)) =
-                                                reply.reply
-                                            {
+                                            if let Some(proto::command_reply::Reply::Text(text)) = reply.reply {
                                                 let dchat_ctx_lock = dchat_ctx.lock().await;
                                                 log::debug!(
                                                     "Command {} :: /{}: editing last message",
                                                     &plugin_name,
                                                     &cmd_name
                                                 );
-                                                if let Err(e) = chat::send_edit_request(
-                                                    &dchat_ctx_lock,
-                                                    msg_id,
-                                                    text,
-                                                )
-                                                .await
+                                                if let Err(e) =
+                                                    chat::send_edit_request(&dchat_ctx_lock, msg_id, text).await
                                                 {
                                                     log::error!("Cannot edit message: {e}");
                                                     continue;
@@ -285,44 +255,27 @@ async fn handle_message(
                                                 &plugin_name
                                             );
                                         }
-                                    } else {
-                                        if let Some(proto::command_reply::Reply::Text(text)) =
-                                            reply.reply
+                                    } else if let Some(proto::command_reply::Reply::Text(text)) = reply.reply {
+                                        let dchat_ctx_lock = dchat_ctx.lock().await;
+                                        log::debug!("Command {} :: /{}: sending message", &plugin_name, &cmd_name);
+                                        let msg_id = match chat::send_text_msg(&dchat_ctx_lock, chat_id, text).await
                                         {
-                                            let dchat_ctx_lock = dchat_ctx.lock().await;
-                                            log::debug!(
-                                                "Command {} :: /{}: sending message",
-                                                &plugin_name,
-                                                &cmd_name
-                                            );
-                                            let msg_id = match chat::send_text_msg(
-                                                &dchat_ctx_lock,
-                                                chat_id,
-                                                text,
-                                            )
-                                            .await
-                                            {
-                                                Ok(msg_id) => msg_id,
-                                                Err(e) => {
-                                                    log::error!("Cannot send new message: {e}");
-                                                    continue;
-                                                }
-                                            };
-                                            last_message_id = Some(msg_id);
-                                        }
+                                            Ok(msg_id) => msg_id,
+                                            Err(e) => {
+                                                log::error!("Cannot send new message: {e}");
+                                                continue;
+                                            }
+                                        };
+                                        last_message_id = Some(msg_id);
                                     }
                                 }
                                 Err(e) => {
-                                    log::error!(
-                                        "Error while executing {} :: /{}: {}",
-                                        &plugin_name,
-                                        &cmd_name,
-                                        e
-                                    );
+                                    log::error!("Error while executing {} :: /{}: {}", &plugin_name, &cmd_name, e);
                                     break;
                                 }
                             }
                         }
+                        log::debug!("{} :: /{} command execution finished", &plugin_name, &cmd_name);
                     }
                     Err(e) => {
                         log::error!(
@@ -374,12 +327,7 @@ async fn main() {
         }
     };
 
-    let requested_plugins: Vec<PluginConfig> = config
-        .plugins
-        .clone()
-        .into_iter()
-        .filter(|p| p.enabled)
-        .collect();
+    let requested_plugins: Vec<PluginConfig> = config.plugins.clone().into_iter().filter(|p| p.enabled).collect();
 
     let bot_context = Arc::new(Mutex::new(BotContext::new(config)));
 

src/paths.rs

@@ -16,11 +16,7 @@ pub(crate) fn default_config_paths() -> Vec<PathBuf> {
         paths.push(config_home.join(APP_CONFIG_DIR).join(CONFIG_FILENAME));
     }
 
-    paths.push(
-        PathBuf::from("/etc")
-            .join(APP_CONFIG_DIR)
-            .join(CONFIG_FILENAME),
-    );
+    paths.push(PathBuf::from("/etc").join(APP_CONFIG_DIR).join(CONFIG_FILENAME));
 
     paths
 }
@@ -28,14 +24,7 @@ pub(crate) fn default_config_paths() -> Vec<PathBuf> {
 pub(crate) fn data_path() -> PathBuf {
     std::env::var("XDG_DATA_HOME")
         .map(|dh| PathBuf::from(dh).join(APP_DATA_DIR))
-        .or_else(|_| {
-            std::env::var("HOME").map(|h| {
-                PathBuf::from(h)
-                    .join(".local")
-                    .join("share")
-                    .join(APP_DATA_DIR)
-            })
-        })
+        .or_else(|_| std::env::var("HOME").map(|h| PathBuf::from(h).join(".local").join("share").join(APP_DATA_DIR)))
         .unwrap_or(PathBuf::from("."))
 }
 

src/plugin/mod.rs

@@ -58,29 +58,23 @@ impl Display for PluginRequestType {
 
 #[async_trait]
 pub(crate) trait PluginConnection: Send + Sync + Debug {
-    async fn initialize_plugin(
-        &self,
-        config: String,
-    ) -> Result<proto::PluginInitializeResponse, PluginConnectionError>;
+    async fn initialize_plugin(&self, config: String)
+    -> Result<proto::PluginInitializeResponse, PluginConnectionError>;
 
-    async fn request_plugin_command_list(
-        &self,
-    ) -> Result<proto::PluginCommandListResponse, PluginConnectionError>;
+    async fn request_plugin_command_list(&self) -> Result<proto::PluginCommandListResponse, PluginConnectionError>;
 
     async fn execute_command(
         self: Arc<Self>,
         command_id: String,
         issuer_id: String,
         argv: Vec<String>,
-    ) -> Result<
-        mpsc::Receiver<Result<proto::CommandReply, PluginConnectionError>>,
-        PluginConnectionError,
-    >;
+    ) -> Result<mpsc::Receiver<Result<proto::CommandReply, PluginConnectionError>>, PluginConnectionError>;
 }
 
 #[derive(Debug)]
 pub(crate) enum PluginConnectionError {
     SendRequest(String, PluginRequestType, Box<dyn Error + Send + Sync>),
+    Timeout(String, PluginRequestType),
     ReadResponse(String, PluginRequestType, Box<dyn Error + Send + Sync>),
     DecodeResponse(String, PluginRequestType, Option<DecodeError>),
     InvalidMessageLength,
@@ -92,6 +86,9 @@ impl Display for PluginConnectionError {
             Self::SendRequest(plugin_name, req_type, e) => f.write_fmt(format_args!(
                 "Can't send request ({req_type}) to plugin {plugin_name}: {e}"
             )),
+            Self::Timeout(plugin_name, req_type) => f.write_fmt(format_args!(
+                "Timed out waiting for response to {req_type} request to plugin {plugin_name}"
+            )),
             Self::ReadResponse(plugin_name, resp_type, e) => f.write_fmt(format_args!(
                 "Can't read response ({resp_type}) from plugin {plugin_name}: {e}"
             )),
@@ -108,6 +105,7 @@ impl Error for PluginConnectionError {
     fn source(&self) -> Option<&(dyn Error + 'static)> {
         match self {
             Self::SendRequest(_, _, e) => Some(e.as_ref()),
+            Self::Timeout(_, _) => None,
             Self::ReadResponse(_, _, e) => Some(e.as_ref()),
             Self::DecodeResponse(_, _, e) => e.as_ref().map(|e| e as &dyn Error),
             Self::InvalidMessageLength => None,
@@ -138,8 +136,7 @@ pub(crate) async fn try_load_plugin(
     cmd.env("DCRCBOT_PLUGIN_TRANSPORT", "stdio");
 
     let plugin_process = cmd.spawn().context("Failed to start the plugin")?;
-    let plugin =
-        stdio::initialize_stdio_plugin(plugin_process, unique_name.clone(), config_json).await?;
+    let plugin = stdio::initialize_stdio_plugin(plugin_process, unique_name.clone(), config_json).await?;
 
     let mut ctx_lock = ctx.lock().await;
     for cmd in plugin.lock().await.commands.iter().cloned() {
@@ -147,12 +144,8 @@ pub(crate) async fn try_load_plugin(
         if ctx_lock.plugin_commands.contains_key(&cmd.name) {
             bail!("duplicate command specification");
         }
-        ctx_lock
-            .plugin_commands
-            .insert(cmd.name.clone(), Arc::clone(&cmd));
-        ctx_lock
-            .plugin_cmd_aliases
-            .insert(cmd.name.clone(), Arc::clone(&cmd));
+        ctx_lock.plugin_commands.insert(cmd.name.clone(), Arc::clone(&cmd));
+        ctx_lock.plugin_cmd_aliases.insert(cmd.name.clone(), Arc::clone(&cmd));
         for alias in cmd.aliases.iter() {
             log::debug!(
                 "adding command alias /{} -> /{} of plugin {}",
@@ -160,9 +153,7 @@ pub(crate) async fn try_load_plugin(
                 &cmd.name,
                 &unique_name
             );
-            ctx_lock
-                .plugin_cmd_aliases
-                .insert(alias.to_owned(), Arc::clone(&cmd));
+            ctx_lock.plugin_cmd_aliases.insert(alias.to_owned(), Arc::clone(&cmd));
         }
     }
 

src/plugin/stdio.rs

@@ -10,9 +10,7 @@ use tokio::{
 };
 
 use crate::{
-    plugin::{
-        LoadedPlugin, PluginCommand, PluginConnection, PluginConnectionError, PluginRequestType,
-    },
+    plugin::{LoadedPlugin, PluginCommand, PluginConnection, PluginConnectionError, PluginRequestType},
     proto,
 };
 
@@ -57,13 +55,19 @@ pub(super) async fn initialize_stdio_plugin(
     Ok(plugin)
 }
 
+#[derive(Debug)]
+enum ResponseListener {
+    Oneshot(oneshot::Sender<proto::Response>),
+    Multishot(mpsc::Sender<proto::Response>),
+}
+
 #[derive(Debug)]
 struct StdioPluginConnection {
     plugin: Arc<Mutex<LoadedPlugin>>,
     process: Mutex<Child>,
     buffered_stdout: Mutex<BufReader<ChildStdout>>,
     next_request_id: Mutex<u32>,
-    pending_requests: Mutex<HashMap<u32, oneshot::Sender<proto::Response>>>,
+    pending_requests: Mutex<HashMap<u32, ResponseListener>>,
 }
 
 impl StdioPluginConnection {
@@ -85,19 +89,16 @@ impl StdioPluginConnection {
 
     async fn stdout_reader_loop(self: Arc<Self>) {
         loop {
-            let frame =
-                match Self::read_length_delimited(self.buffered_stdout.lock().await.deref_mut())
-                    .await
-                {
-                    Ok(frame) => frame,
-                    Err(e) => {
-                        log::error!(
-                            "Error while reading STDOUT of stdio plugin {}: {e}",
-                            &self.plugin.lock().await.plugin_id
-                        );
-                        break;
-                    }
-                };
+            let frame = match Self::read_length_delimited(self.buffered_stdout.lock().await.deref_mut()).await {
+                Ok(frame) => frame,
+                Err(e) => {
+                    log::error!(
+                        "Error while reading STDOUT of stdio plugin {}: {e}",
+                        &self.plugin.lock().await.plugin_id
+                    );
+                    break;
+                }
+            };
             let response = match proto::Response::decode(frame.as_slice()) {
                 Ok(response) => response,
                 Err(_) => {
@@ -108,18 +109,29 @@ impl StdioPluginConnection {
                     continue;
                 }
             };
-            match self
-                .pending_requests
-                .lock()
-                .await
-                .remove(&response.request_id)
-            {
-                Some(sender) => match sender.send(response) {
+            let mut pending_req_lock = self.pending_requests.lock().await;
+            match pending_req_lock.get(&response.request_id) {
+                Some(ResponseListener::Oneshot(_)) => {
+                    // We must remove to be able to send the response to the oneshot channel
+                    if let Some(ResponseListener::Oneshot(sender)) = pending_req_lock.remove(&response.request_id) {
+                        match sender.send(response) {
+                            Ok(()) => {}
+                            Err(response) => {
+                                log::warn!(
+                                    "Dropping response with request_id {} from plugin {}",
+                                    response.request_id,
+                                    self.plugin.lock().await.plugin_id
+                                );
+                            }
+                        }
+                    }
+                }
+                Some(ResponseListener::Multishot(sender)) => match sender.send(response).await {
                     Ok(()) => {}
-                    Err(response) => {
+                    Err(e) => {
                         log::warn!(
                             "Dropping response with request_id {} from plugin {}",
-                            response.request_id,
+                            e.0.request_id,
                             self.plugin.lock().await.plugin_id
                         );
                     }
@@ -133,9 +145,7 @@ impl StdioPluginConnection {
         self.pending_requests.lock().await.clear();
     }
 
-    async fn read_length_delimited<R>(
-        reader: &mut R,
-    ) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>>
+    async fn read_length_delimited<R>(reader: &mut R) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>>
     where
         R: AsyncRead + Unpin,
     {
@@ -162,21 +172,6 @@ impl StdioPluginConnection {
         reader.read_exact(&mut message).await?;
         Ok(message)
     }
-
-    async fn await_response_to(
-        &self,
-        request_id: u32,
-        timeout: Duration,
-    ) -> Result<proto::Response, Box<dyn Error + Send + Sync>> {
-        let (tx, rx) = oneshot::channel();
-        self.pending_requests.lock().await.insert(request_id, tx);
-
-        match tokio::time::timeout(timeout, rx).await {
-            Ok(Ok(response)) => Ok(response),
-            Ok(Err(e)) => Err(Box::new(e)),
-            Err(e) => Err(Box::new(e)),
-        }
-    }
 }
 
 #[async_trait]
@@ -185,20 +180,31 @@ impl PluginConnection for StdioPluginConnection {
         &self,
         config: String,
     ) -> Result<proto::PluginInitializeResponse, PluginConnectionError> {
+        // Generate request id
         let request_id = {
             let mut r = self.next_request_id.lock().await;
             let id = *r;
             *r += 1;
             id
         };
+
+        // Encode request
         let request = proto::Request {
             request_id,
-            req: Some(proto::request::Req::InitializeReq(
-                proto::PluginInitializeRequest { config },
-            )),
+            req: Some(proto::request::Req::InitializeReq(proto::PluginInitializeRequest {
+                config,
+            })),
         }
         .encode_length_delimited_to_vec();
 
+        // Add response listener
+        let (tx, rx) = oneshot::channel();
+        self.pending_requests
+            .lock()
+            .await
+            .insert(request_id, ResponseListener::Oneshot(tx));
+
+        // Send request
         if let Err(e) = self
             .process
             .lock()
@@ -216,16 +222,20 @@ impl PluginConnection for StdioPluginConnection {
             ));
         }
 
-        let response = match self
-            .await_response_to(request_id, Duration::from_secs(10))
-            .await
-        {
-            Ok(response) => response,
-            Err(e) => {
+        // Receive response
+        let response = match tokio::time::timeout(Duration::from_secs(10), rx).await {
+            Ok(Ok(response)) => response,
+            Ok(Err(e)) => {
                 return Err(PluginConnectionError::ReadResponse(
                     self.plugin.lock().await.plugin_id.clone(),
                     PluginRequestType::Initialize,
-                    e,
+                    Box::new(e),
+                ));
+            }
+            Err(_) => {
+                return Err(PluginConnectionError::Timeout(
+                    self.plugin.lock().await.plugin_id.clone(),
+                    PluginRequestType::Initialize,
                 ));
             }
         };
@@ -240,23 +250,30 @@ impl PluginConnection for StdioPluginConnection {
         }
     }
 
-    async fn request_plugin_command_list(
-        &self,
-    ) -> Result<proto::PluginCommandListResponse, PluginConnectionError> {
+    async fn request_plugin_command_list(&self) -> Result<proto::PluginCommandListResponse, PluginConnectionError> {
+        // Generate request id
         let request_id = {
             let mut r = self.next_request_id.lock().await;
             let id = *r;
             *r += 1;
             id
         };
+
+        // Encode request
         let request = proto::Request {
             request_id,
-            req: Some(proto::request::Req::CommandListReq(
-                proto::PluginCommandListRequest {},
-            )),
+            req: Some(proto::request::Req::CommandListReq(proto::PluginCommandListRequest {})),
         }
         .encode_length_delimited_to_vec();
 
+        // Add response listener
+        let (tx, rx) = oneshot::channel();
+        self.pending_requests
+            .lock()
+            .await
+            .insert(request_id, ResponseListener::Oneshot(tx));
+
+        // Send request
         if let Err(e) = self
             .process
             .lock()
@@ -274,16 +291,20 @@ impl PluginConnection for StdioPluginConnection {
             ));
         }
 
-        let response = match self
-            .await_response_to(request_id, Duration::from_secs(10))
-            .await
-        {
-            Ok(response) => response,
-            Err(e) => {
+        // Receive response
+        let response = match tokio::time::timeout(Duration::from_secs(10), rx).await {
+            Ok(Ok(response)) => response,
+            Ok(Err(e)) => {
                 return Err(PluginConnectionError::ReadResponse(
                     self.plugin.lock().await.plugin_id.clone(),
                     PluginRequestType::CommandList,
-                    e,
+                    Box::new(e),
+                ));
+            }
+            Err(_) => {
+                return Err(PluginConnectionError::Timeout(
+                    self.plugin.lock().await.plugin_id.clone(),
+                    PluginRequestType::CommandList,
                 ));
             }
         };
@@ -303,28 +324,34 @@ impl PluginConnection for StdioPluginConnection {
         command_id: String,
         issuer_id: String,
         argv: Vec<String>,
-    ) -> Result<
-        mpsc::Receiver<Result<proto::CommandReply, PluginConnectionError>>,
-        PluginConnectionError,
-    > {
+    ) -> Result<mpsc::Receiver<Result<proto::CommandReply, PluginConnectionError>>, PluginConnectionError> {
+        // Generate request id
         let request_id = {
             let mut r = self.next_request_id.lock().await;
             let id = *r;
             *r += 1;
             id
         };
+
+        // Encode request
         let request = proto::Request {
             request_id,
-            req: Some(proto::request::Req::ExecuteReq(
-                proto::PluginExecuteRequest {
-                    command_id,
-                    issuer_id,
-                    arg_vector: argv,
-                },
-            )),
+            req: Some(proto::request::Req::ExecuteReq(proto::PluginExecuteRequest {
+                command_id,
+                issuer_id,
+                arg_vector: argv,
+            })),
         }
         .encode_length_delimited_to_vec();
 
+        // Add response listener
+        let (tx, mut rx) = mpsc::channel(4);
+        self.pending_requests
+            .lock()
+            .await
+            .insert(request_id, ResponseListener::Multishot(tx));
+
+        // Send request
         if let Err(e) = self
             .process
             .lock()
@@ -342,26 +369,14 @@ impl PluginConnection for StdioPluginConnection {
             ));
         }
 
-        let (tx, rx) = mpsc::channel(4);
-
+        // Receive responses
+        let (tx_out, rx_out) = mpsc::channel(4);
         tokio::spawn(async move {
             loop {
-                let response = match self
-                    .await_response_to(request_id, Duration::from_secs(600))
-                    .await
-                {
-                    Ok(response) => response,
-                    Err(e) => {
-                        if let Err(e) = tx
-                            .send(Err(PluginConnectionError::ReadResponse(
-                                self.plugin.lock().await.plugin_id.clone(),
-                                PluginRequestType::Execute,
-                                e,
-                            )))
-                            .await
-                        {
-                            log::error!("Cannot send error notification to another task: {e}");
-                        }
+                let response = match rx.recv().await {
+                    Some(response) => response,
+                    None => {
+                        log::error!("Response channel for request {request_id} closed unexpectedly");
                         break;
                     }
                 };
@@ -372,7 +387,8 @@ impl PluginConnection for StdioPluginConnection {
                             Some(proto::command_reply::Reply::End(_)) => true,
                             _ => false,
                         };
-                        if let Err(e) = tx.send(Ok(reply)).await {
+                        log::trace!("reply sent: Ok({:?})", &reply);
+                        if let Err(e) = tx_out.send(Ok(reply)).await {
                             log::error!("Cannot send command reply to another task: {e}");
                         }
                         if end {
@@ -380,7 +396,7 @@ impl PluginConnection for StdioPluginConnection {
                         }
                     }
                     _ => {
-                        if let Err(e) = tx
+                        if let Err(e) = tx_out
                             .send(Err(PluginConnectionError::DecodeResponse(
                                 self.plugin.lock().await.plugin_id.clone(),
                                 PluginRequestType::Execute,
@@ -394,8 +410,10 @@ impl PluginConnection for StdioPluginConnection {
                     }
                 }
             }
+            // Remove response listener
+            self.pending_requests.lock().await.remove_entry(&request_id);
         });
 
-        Ok(rx)
+        Ok(rx_out)
     }
 }

src/ssh.rs

@@ -36,11 +36,7 @@ pub(crate) async fn connect_ssh<A: ToSocketAddrs + Debug, K: AsRef<[u8]>, P: Int
         "Connecting to SSH {:?} {} key & {} password",
         &addr,
         if key.is_some() { "with" } else { "without" },
-        if password.is_some() {
-            "with"
-        } else {
-            "without"
-        }
+        if password.is_some() { "with" } else { "without" }
     );
     let mut connection = tokio::time::timeout(
         connect_timeout,
@@ -65,12 +61,7 @@ pub(crate) async fn connect_ssh<A: ToSocketAddrs + Debug, K: AsRef<[u8]>, P: Int
     if !auth_success {
         if let Some(key) = key {
             log::info!("Trying to authenticate using key");
-            match tokio::time::timeout(
-                auth_timeout,
-                connection.authenticate_publickey(username, key),
-            )
-            .await
-            {
+            match tokio::time::timeout(auth_timeout, connection.authenticate_publickey(username, key)).await {
                 Ok(Ok(AuthResult::Success)) => {
                     auth_success = true;
                 }
@@ -78,8 +69,7 @@ pub(crate) async fn connect_ssh<A: ToSocketAddrs + Debug, K: AsRef<[u8]>, P: Int
                     remaining_methods,
                     partial_success,
                 })) => {
-                    if partial_success || !remaining_methods.contains(&russh::MethodKind::Password)
-                    {
+                    if partial_success || !remaining_methods.contains(&russh::MethodKind::Password) {
                         if partial_success {
                             anyhow::bail!("SSH auth failed: multi-factor auth is not supported");
                         } else {
@@ -102,12 +92,7 @@ pub(crate) async fn connect_ssh<A: ToSocketAddrs + Debug, K: AsRef<[u8]>, P: Int
     if !auth_success {
         if let Some(password) = password {
             log::info!("Trying to authenticate using password");
-            match tokio::time::timeout(
-                auth_timeout,
-                connection.authenticate_password(username, password),
-            )
-            .await
-            {
+            match tokio::time::timeout(auth_timeout, connection.authenticate_password(username, password)).await {
                 Ok(Ok(AuthResult::Success)) => {
                     auth_success = true;
                 }