rift: slightly better error handling

1 files changed, 123 insertions, 112 deletions

diff --git a/client-native-rift/src/main.rs b/client-native-rift/src/main.rs
index c44f8c2..fad4bb7 100644
--- a/client-native-rift/src/main.rs
+++ b/client-native-rift/src/main.rs
@@ -7,6 +7,7 @@
 pub mod file;
 pub mod port;
 
+use anyhow::bail;
 use clap::{ColorChoice, Parser};
 use file::{DownloadHandler, FileSender};
 use libkeks::{
@@ -21,7 +22,7 @@ use port::{ForwardHandler, PortExposer};
 use rustyline::{error::ReadlineError, DefaultEditor};
 use std::{
     collections::HashMap, future::Future, os::unix::prelude::MetadataExt, path::PathBuf, pin::Pin,
-    sync::Arc,
+    process::exit, sync::Arc,
 };
 use tokio::{fs, net::TcpListener, sync::RwLock};
 use users::get_current_username;
@@ -40,8 +41,9 @@ fn main() {
         .unwrap();
 }
 
-#[derive(Parser, Clone)]
 /// If no command is provided, rift will enter REPL-mode.
+#[derive(Parser, Clone)]
+#[clap(multicall = false, color = ColorChoice::Auto)]
 pub struct Args {
     /// keks-meet server used for establishing p2p connection
     #[clap(long, default_value = "wss://meet.metamuffin.org")]
@@ -52,8 +54,8 @@ pub struct Args {
     /// pre-shared secret (aka. room name)
     secret: String,
     // /// Dispatch a single command after startup
-    // #[clap(subcommand)]
-    // command: Option<Command>,
+    #[clap(subcommand)]
+    command: Option<Command>,
 }
 
 #[derive(Parser, Debug, Clone)]
@@ -115,118 +117,21 @@ async fn run() -> anyhow::Result<()> {
     inst.spawn_ping().await;
     tokio::task::spawn(inst.clone().receive_loop());
 
+    if let Some(command) = args.command {
+        info!("running startup command...");
+        if let Err(e) = dispatch_command(&inst, &state, command).await {
+            error!("{e}");
+            exit(1);
+        };
+        info!("done");
+    }
     let mut rl = DefaultEditor::new()?;
     loop {
         match rl.readline("> ") {
             Ok(line) => match Command::try_parse_from(shlex::split(&line).unwrap()) {
-                Ok(command) => match command {
-                    Command::List => {
-                        let peers = inst.peers.read().await;
-                        println!("{} clients available", peers.len());
-                        for p in peers.values() {
-                            let username = p
-                                .username
-                                .read()
-                                .await
-                                .clone()
-                                .unwrap_or("<unknown>".to_string());
-                            println!("{username}:");
-                            for (rid, r) in p.remote_provided.read().await.iter() {
-                                println!(
-                                    "\t{rid:?}: {} {:?}",
-                                    r.kind,
-                                    r.label.clone().unwrap_or_default()
-                                )
-                            }
-                        }
-                    }
-                    Command::Provide { path, id } => {
-                        inst.add_local_resource(Box::new(FileSender {
-                            info: ProvideInfo {
-                                id: id.unwrap_or("file".to_owned()),
-                                kind: "file".to_string(),
-                                track_kind: None,
-                                label: Some(
-                                    path.file_name().unwrap().to_str().unwrap().to_string(),
-                                ),
-                                size: Some(fs::metadata(&path).await.unwrap().size() as usize),
-                            },
-                            path: path.into(),
-                        }))
-                        .await;
-                    }
-                    Command::Download { id, path } => {
-                        let peers = inst.peers.read().await;
-                        'outer: for p in peers.values() {
-                            for (rid, r) in p.remote_provided.read().await.iter() {
-                                if rid == &id {
-                                    if r.kind == "file" {
-                                        state
-                                            .write()
-                                            .await
-                                            .requested
-                                            .insert(id.clone(), Box::new(DownloadHandler { path }));
-                                        p.request_resource(id).await;
-                                    } else {
-                                        warn!("not a file");
-                                    }
-                                    break 'outer;
-                                }
-                            }
-                        }
-                    }
-                    Command::Expose { port, id } => {
-                        inst.add_local_resource(Box::new(PortExposer {
-                            port,
-                            info: ProvideInfo {
-                                kind: "port".to_string(),
-                                id: id.unwrap_or(format!("p{port}")),
-                                track_kind: None,
-                                label: Some(format!("port {port}")),
-                                size: None,
-                            },
-                        }))
-                        .await;
-                    }
-                    Command::Forward { id, port } => {
-                        let peers = inst.peers.read().await;
-                        'outer: for peer in peers.values() {
-                            for (rid, r) in peer.remote_provided.read().await.iter() {
-                                if rid == &id {
-                                    if r.kind == "port" {
-                                        let peer = peer.to_owned();
-                                        let state = state.clone();
-                                        tokio::task::spawn(async move {
-                                            let Ok(listener) =
-                                                TcpListener::bind(("127.0.0.1", port.unwrap_or(0)))
-                                                    .await
-                                            else {
-                                                error!("cannot bind tcp listener");
-                                                return;
-                                            };
-                                            info!(
-                                                "tcp listener bound to {}",
-                                                listener.local_addr().unwrap()
-                                            );
-                                            while let Ok((stream, addr)) = listener.accept().await {
-                                                info!("new connection from {addr:?}");
-                                                state.write().await.requested.insert(
-                                                    id.clone(),
-                                                    Box::new(ForwardHandler {
-                                                        stream: Arc::new(RwLock::new(Some(stream))),
-                                                    }),
-                                                );
-                                                peer.request_resource(id.clone()).await;
-                                            }
-                                        });
-                                    } else {
-                                        warn!("not a port");
-                                    }
-                                    break 'outer;
-                                }
-                            }
-                        }
-                    }
+                Ok(command) => match dispatch_command(&inst, &state, command).await {
+                    Ok(()) => (),
+                    Err(err) => error!("{err}"),
                 },
                 Err(err) => err.print().unwrap(),
             },
@@ -245,6 +150,112 @@ async fn run() -> anyhow::Result<()> {
     Ok(())
 }
 
+async fn dispatch_command(
+    inst: &Arc<Instance>,
+    state: &Arc<RwLock<State>>,
+    command: Command,
+) -> anyhow::Result<()> {
+    match command {
+        Command::List => {
+            let peers = inst.peers.read().await;
+            println!("{} clients available", peers.len());
+            for p in peers.values() {
+                let username = p
+                    .username
+                    .read()
+                    .await
+                    .clone()
+                    .unwrap_or("<unknown>".to_string());
+                println!("{username}:");
+                for (rid, r) in p.remote_provided.read().await.iter() {
+                    println!(
+                        "\t{rid:?}: {} {:?}",
+                        r.kind,
+                        r.label.clone().unwrap_or_default()
+                    )
+                }
+            }
+        }
+        Command::Provide { path, id } => {
+            inst.add_local_resource(Box::new(FileSender {
+                info: ProvideInfo {
+                    id: id.unwrap_or("file".to_owned()),
+                    kind: "file".to_string(),
+                    track_kind: None,
+                    label: Some(path.file_name().unwrap().to_str().unwrap().to_string()),
+                    size: Some(fs::metadata(&path).await.unwrap().size() as usize),
+                },
+                path: path.into(),
+            }))
+            .await;
+        }
+        Command::Download { id, path } => {
+            let (peer, _resource) = find_id(inst, id.clone(), "file").await?;
+            state
+                .write()
+                .await
+                .requested
+                .insert(id.clone(), Box::new(DownloadHandler { path }));
+            peer.request_resource(id).await;
+        }
+        Command::Expose { port, id } => {
+            inst.add_local_resource(Box::new(PortExposer {
+                port,
+                info: ProvideInfo {
+                    kind: "port".to_string(),
+                    id: id.unwrap_or(format!("p{port}")),
+                    track_kind: None,
+                    label: Some(format!("port {port}")),
+                    size: None,
+                },
+            }))
+            .await;
+        }
+        Command::Forward { id, port } => {
+            let (peer, _resource) = find_id(inst, id.clone(), "port").await?;
+            let state = state.clone();
+            tokio::task::spawn(async move {
+                let Ok(listener) = TcpListener::bind(("127.0.0.1", port.unwrap_or(0))).await else {
+                    error!("cannot bind tcp listener");
+                    return;
+                };
+                info!("tcp listener bound to {}", listener.local_addr().unwrap());
+                while let Ok((stream, addr)) = listener.accept().await {
+                    info!("new connection from {addr:?}");
+                    state.write().await.requested.insert(
+                        id.clone(),
+                        Box::new(ForwardHandler {
+                            stream: Arc::new(RwLock::new(Some(stream))),
+                        }),
+                    );
+                    peer.request_resource(id.clone()).await;
+                }
+            });
+        }
+    }
+    Ok(())
+}
+
+async fn find_id(
+    inst: &Arc<Instance>,
+    id: String,
+    kind: &str,
+) -> anyhow::Result<(Arc<Peer>, ProvideInfo)> {
+    let peers = inst.peers.read().await;
+    for peer in peers.values() {
+        for (rid, r) in peer.remote_provided.read().await.iter() {
+            if rid == &id {
+                if r.kind == kind {
+                    return Ok((peer.to_owned(), r.to_owned()));
+                } else {
+                    bail!("wrong type: expected {kind:?}, found {:?}", r.kind)
+                }
+            }
+        }
+    }
+    bail!("id not found")
+}
+
 impl EventHandler for Handler {
     fn peer_join(&self, _peer: Arc<Peer>) -> libkeks::DynFut<()> {
         Box::pin(async move {})