yodle/src/server.rs

// Copyright (c) 2019 Reyk Floeter <contact@reykfloeter.com>
//
// Permission to use, copy, modify, and distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice appear in all copies.
//
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
// WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
// MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
// ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
// WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
// ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
// OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

use crate::Config;
use dialog::{backends::Zenity, DialogBox};
use dirs::home_dir;
use std::{
    io::{Error, ErrorKind, Result},
    path::{PathBuf, MAIN_SEPARATOR},
};
use tokio::{
    fs::{remove_file, File},
    io::{
        copy, split, AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWriteExt, BufReader, BufWriter,
    },
    net::TcpListener,
};
use tokio_libtls::prelude::*;

const MAX_COPIES: usize = 1024;
const DOWNLOAD_DIR: &str = "Downloads";

pub(crate) async fn run(config: Config) -> Result<()> {
    let (cert, key, ca) = config.load_keys()?;
    let mut options = config.load_server_options();
    let home = home_dir().unwrap();
    let mut download_dir = PathBuf::from(home);

    download_dir.push(DOWNLOAD_DIR);

    let tls_config = TlsConfigBuilder::new()
        .ca_file(ca)
        .keypair_file(cert, key, None)
        .build()
        .unwrap();

    let addr = config.address.unwrap();
    let mut listener = TcpListener::bind(&addr).await?;

    loop {
        let (tcp, _) = listener.accept().await?;
        let size_limit = config.size_limit as u64;
        let peer_addr: String = tcp.peer_addr()?.to_string();
        let dont_ask = config.dont_ask;
        let options = options.build();
        let tls = match AsyncTls::accept_stream(tcp, &tls_config, options).await {
            Ok(tls) => {
                debug!("{} status: connected", peer_addr);
                tls
            }
            Err(err) => {
                debug!("{} status: TLS connection error ({})", peer_addr, err);
                continue;
            }
        };

        let (reader, mut writer) = split(tls);
        let mut reader = BufReader::new(reader);
        let mut filename = download_dir.clone();

        tokio::spawn(async move {
            // Read and validate the filename
            let line = match read_line(&peer_addr, &mut reader).await {
                Ok(s) if !(s.contains(MAIN_SEPARATOR) || s.contains('/') || s.contains('\\')) => s,
                Err(err) => {
                    debug!("{}", err);
                    return;
                }
                _ => {
                    debug!("{} failed: filename invalid", peer_addr);
                    return;
                }
            };
            filename.push(&line);

            // If the filename exists, try to append a counter (e.g. foo_1.gz)
            for i in 1..=(MAX_COPIES + 1) {
                if i == MAX_COPIES {
                    debug!(
                        "{} failed: {} exists more than {} times",
                        peer_addr, line, i
                    );
                    return;
                }
                if !filename.exists() {
                    break;
                }
                let ext = line.rfind('.').unwrap_or(line.len());
                filename.set_file_name(format!("{}_{}{}", &line[..ext], i, &line[ext..]));
            }

            // Read the file size
            let line = match read_line(&peer_addr, &mut reader).await {
                Ok(s) => s,
                Err(err) => {
                    debug!("{}", err);
                    return;
                }
            };
            let file_size: u64 = match line.parse() {
                Ok(s) => s,
                Err(err) => {
                    debug!("{} failed: file size ({})", peer_addr, err);
                    return;
                }
            };
            if file_size == 0 || (size_limit > 0 && file_size > size_limit) {
                debug!(
                    "{} failed: file size (out of limits, {} bytes)",
                    peer_addr, file_size
                );
                return;
            }

            // Print dialog
            if !dont_ask {
                let mut zenity = Zenity::new();
                zenity.set_icon("question");
                zenity.set_width(360);
                let choice = dialog::Question::new(&format!(
                    "Do you want to accept\n{}\n({} bytes)?",
                    filename.file_name().as_ref().unwrap().to_string_lossy(),
                    file_size
                ))
                .title("Yodle!")
                .show_with(&zenity)
                .expect("Could not display dialog box");
                if choice != dialog::Choice::Yes {
                    info!("{} failed: rejected {}", peer_addr, filename.display(),);
                    return;
                }
            }

            debug!(
                "{} status: receiving {} ({} bytes)",
                peer_addr,
                filename.display(),
                file_size
            );

            // Create output file
            let file = match File::create(&filename).await {
                Ok(f) => f,
                Err(err) => {
                    debug!(
                        "{} failed {}: file ({})",
                        peer_addr,
                        filename.display(),
                        err
                    );
                    return;
                }
            };

            // I/O
            let mut file_writer = BufWriter::new(file);
            let mut reader = reader.take(file_size);
            let copied = match copy(&mut reader, &mut file_writer).await {
                Ok(s) => s,
                Err(err) => {
                    debug!("{} failed: I/O ({})", peer_addr, err);
                    return;
                }
            };

            // Check result and send response
            let result = if copied != file_size {
                drop(file_writer);
                let _ = remove_file(&filename).await.is_ok();
                warn!(
                    "{} failed: {} ({}/{} bytes)",
                    peer_addr,
                    filename.display(),
                    copied,
                    file_size
                );
                "failed: truncated file\n".to_string()
            } else {
                info!(
                    "{} success: {} ({} bytes)",
                    peer_addr,
                    filename.display(),
                    copied
                );
                "success\n".to_string()
            };

            let _ = writer.write_all(result.as_bytes()).await;
        });
    }
}

pub async fn read_line<T: AsyncRead + AsyncBufReadExt + Unpin>(
    peer: &str,
    reader: &mut T,
) -> Result<String> {
    let mut line = String::with_capacity(1024);
    let mut len = 0;

    // Ignore some empty lines
    for _ in 0..10 {
        if let Err(err) = reader.read_line(&mut line).await {
            return Err(Error::new(
                ErrorKind::Other,
                format!("{} failed: read ({})", peer, err),
            ));
        }
        len = line.find(|c: char| c == '\r' || c == '\n').unwrap_or(0);
        if len > 0 {
            break;
        }
    }

    Ok((&line[0..len]).to_owned())
}