quick_file_transfer/
server.rs

1use crate::{
2    config::{
3        transfer::{
4            command::{ServerCommand, ServerResult},
5            listen::ListenArgs,
6        },
7        Config,
8    },
9    util::{read_server_cmd, server_handshake},
10};
11use anyhow::{bail, Result};
12use std::{
13    net::{IpAddr, TcpListener},
14    path::PathBuf,
15    sync::{
16        atomic::{AtomicBool, Ordering},
17        Arc,
18    },
19};
20
21mod path;
22use path::validate_remote_path;
23
24pub mod util;
25use util::{join_all_threads, send_result, spawn_child_on_new_port};
26
27pub mod child;
28
29pub fn listen(_cfg: &Config, listen_args: &ListenArgs) -> Result<()> {
30    let ListenArgs {
31        ip,
32        port,
33        output: _,
34        decompression: _,
35        output_dir: _,
36        remote: _,
37    } = listen_args;
38
39    let stop_flag: Arc<AtomicBool> = Arc::new(AtomicBool::new(false));
40    let ip: IpAddr = ip.parse()?;
41    let initial_listener = TcpListener::bind((ip, *port))?;
42    run_server(&initial_listener, listen_args, &stop_flag)
43}
44
45fn run_server(
46    initial_listener: &TcpListener,
47    args: &ListenArgs,
48    stop_flag: &Arc<AtomicBool>,
49) -> anyhow::Result<()> {
50    let mut thread_handles = vec![];
51    match initial_listener.accept() {
52        Ok((mut socket, addr)) => {
53            tracing::info!("Client accepted at: {addr:?}");
54            server_handshake(&mut socket)?;
55            let mut root_dest: Option<PathBuf> = None; // Used as root destination if invoked through ssh/scp mode
56            let mut cmd_buf: [u8; 256] = [0; 256];
57            loop {
58                if let Some(cmd) = read_server_cmd(&mut socket, &mut cmd_buf)? {
59                    tracing::trace!("Received command: {cmd:?}");
60                    match cmd {
61                        ServerCommand::GetFreePort(_) => {
62                            let child_thread_handle = spawn_child_on_new_port(
63                                &mut socket,
64                                args,
65                                &Arc::clone(stop_flag),
66                                &cmd,
67                                root_dest.clone(),
68                            )?;
69                            thread_handles.push(child_thread_handle);
70                        }
71                        ServerCommand::EndOfTransfer => {
72                            tracing::trace!("Received command: {cmd:?}, stopping all threads...");
73                            stop_flag.store(true, Ordering::Relaxed);
74                            match join_all_threads(thread_handles) {
75                                Ok(_) => {
76                                    send_result(&mut socket, &ServerResult::Ok)?;
77                                    return Ok(());
78                                }
79                                Err(th_errs) => {
80                                    let err_res = ServerResult::err(th_errs.clone());
81                                    send_result(&mut socket, &err_res)?;
82                                    bail!(th_errs);
83                                }
84                            }
85                        }
86                        ServerCommand::IsDestinationValid(mode, dest) => {
87                            let dest = PathBuf::from(dest);
88                            tracing::info!("Checking validity of remote path: {dest:?}");
89                            match validate_remote_path(&mode, &dest) {
90                                Ok(remote_dest) => {
91                                    send_result(&mut socket, &ServerResult::Ok)?;
92                                    root_dest = Some(remote_dest);
93                                }
94                                Err(e) => {
95                                    tracing::error!("Invalid remote path: {e}");
96                                    send_result(
97                                        &mut socket,
98                                        &ServerResult::Err(e.to_string().into()),
99                                    )?;
100                                }
101                            }
102                        }
103                        // For child threads
104                        ServerCommand::Prealloc(_, _) => todo!(),
105                        ServerCommand::ReceiveData(_, _, _) => todo!(),
106                    }
107                } else {
108                    tracing::debug!("Main Client disconnected...");
109                    break;
110                }
111            }
112        }
113        Err(e) => bail!(e),
114    }
115    Ok(())
116}
quick_file_transfer/server.rs

quick_file_transfer/
server.rs
