227 lines
7.5 KiB
Rust
227 lines
7.5 KiB
Rust
// 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())
|
|
}
|