implement server part

Cargo.toml

@@ -5,6 +5,7 @@ authors = ["Reyk Floeter <contact@reykfloeter.com>"]
 edition = "2018"
 
 [dependencies]
+bubblebabble = "0.1"
 derive_more = "0.99"
 dirs = "2.0"
 env_logger = "0.7"

src/cert.rs

@@ -12,11 +12,14 @@
 // 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 bubblebabble::stablebabble;
 use derive_more::{Deref, From, Into};
 use dirs::home_dir;
 use std::{
     fs::{create_dir, File},
     io::{Error, ErrorKind, Result, Write},
+    net::SocketAddr,
     path::PathBuf,
     process::{Command, Output},
 };
@@ -37,7 +40,7 @@ pub struct Cert {
 
 /// The local server keypair
 #[derive(Clone, Debug, Deref, From, Into)]
-pub struct KeyPair(Cert);
+pub(crate) struct KeyPair(Cert);
 
 impl KeyPair {
     /// Get the server keys or generate them if they don't exist
@@ -45,8 +48,9 @@ impl KeyPair {
     /// # Panics
     ///
     /// This function may fail if the key cannot be generated.
-    pub fn new() -> Self {
+    pub(crate) fn new(config: &Config) -> Self {
         let home = home_dir().unwrap();
+
         let mut config_dir = PathBuf::from(home);
         config_dir.push(CONFIG_DIR);
 
@@ -73,7 +77,7 @@ impl KeyPair {
         }
 
         if !csr.exists() {
-            gen_csr(&key_path, &csr_path).expect("failed to generate signing request");
+            gen_csr(&key_path, &csr_path, config).expect("failed to generate signing request");
         }
 
         if !cert.exists() {
@@ -108,13 +112,31 @@ fn gen_key(key_path: &str) -> Result<()> {
     Ok(())
 }
 
-fn gen_csr(key_path: &str, csr_path: &str) -> Result<()> {
+fn gen_csr(key_path: &str, csr_path: &str, config: &Config) -> Result<()> {
     debug!("key {} csr {}", key_path, csr_path);
 
     // XXX use gethostname() libc function or a crate around it
     let hostname = String::from_utf8(Command::new("hostname").output()?.stdout)
         .map_err(|err| Error::new(ErrorKind::Other, err))?;
-    let subject = format!("/C={cn}/CN={hn}/O=yodle/", cn = COUNTRY, hn = hostname);
+    let hostname = hostname.trim();
+
+    let babble = if let Some(SocketAddr::V6(ref address)) = config.address {
+        stablebabble(&address.ip().octets())
+    } else {
+        return Err(Error::new(ErrorKind::Other, "invalid address"));
+    };
+
+    let subject = format!(
+        "/C={cn}/CN={hn}/O={sn}/OU={bb}",
+        cn = COUNTRY,
+        hn = hostname,
+        sn = config
+            .servername
+            .as_ref()
+            .map(|s| s.as_str())
+            .unwrap_or("yodle"),
+        bb = babble,
+    );
 
     command_ok(
         Command::new("openssl")
@@ -126,7 +148,16 @@ fn gen_csr(key_path: &str, csr_path: &str) -> Result<()> {
     Command::new("chmod").args(&["0640", csr_path]).output()?;
 
     let mut f = File::create(csr_path.to_owned() + ".ext")?;
-    f.write_all(format!("subjectAltName=DNS:localhost,DNS:{}", hostname).as_bytes())?;
+    let mut san = format!(
+        "subjectAltName=DNS:localhost,DNS:{hn},DNS:{bb}",
+        hn = hostname,
+        bb = babble
+    );
+    if let Some(ref servername) = config.servername {
+        san.push_str(&format!(",DNS:{}", servername));
+    }
+    eprintln!("{}", san);
+    f.write_all(san.as_bytes())?;
     Ok(())
 }
 

src/main.rs

@@ -19,12 +19,16 @@ mod cert;
 mod client;
 mod server;
 
+use bubblebabble::stablebabble;
 use cert::KeyPair;
+use getopts::Options;
 use log::LevelFilter;
 use std::{
+    env,
     io::{Error, ErrorKind, Result},
     net::SocketAddr,
     path::{Path, PathBuf},
+    process,
     time::Duration,
 };
 use tokio_libtls::prelude::*;
@@ -36,12 +40,16 @@ pub(crate) struct Config {
     timeout: Option<Duration>,
     servername: Option<String>,
     address: Option<SocketAddr>,
+    size_limit: usize,
 }
 
 impl Config {
     pub fn new() -> Self {
         Self {
             address: "[::1]:8023".parse().ok(),
+            servername: Some(env::var("USER").unwrap_or("localhost".to_string())),
+            timeout: Some(Duration::from_secs(10)),
+            size_limit: 1_073_741_824,
             ..Default::default()
         }
     }
@@ -77,19 +85,64 @@ impl Config {
     }
 }
 
+fn usage(program: &str, opts: Options) -> ! {
+    let brief = format!("Usage: {} [options]", program);
+    print!("{}", opts.usage(&brief));
+    process::exit(1)
+}
+
 #[tokio::main]
 async fn main() {
+    let args: Vec<String> = env::args().collect();
+    let program = args[0].clone();
     let mut config = Config::new();
 
+    let mut opts = Options::new();
+    opts.optopt(
+        "a",
+        "address",
+        "client/server address",
+        &config.address.unwrap().to_string(),
+    );
+    opts.optopt(
+        "n",
+        "name",
+        "server name",
+        &config.servername.as_ref().unwrap(),
+    );
+    opts.optflag("s", "server", "run server");
+    opts.optflag("c", "client", "connect as client");
+    opts.optflag("h", "help", "print this help menu");
+    let matches = match opts.parse(&args[1..]) {
+        Ok(m) => m,
+        Err(f) => panic!(f.to_string()),
+    };
+    if matches.opt_present("h") || (matches.opt_present("c") && matches.opt_present("s")) {
+        usage(&program, opts);
+    }
+
     env_logger::builder()
         .filter_level(LevelFilter::Debug)
         .init();
 
-    let keypair = KeyPair::new();
+    let addr = match config.address {
+        Some(SocketAddr::V6(addr)) => addr.clone(),
+        _ => panic!("invalid address: {:?}", config.address),
+    };
+
+    let keypair = KeyPair::new(&config);
     config.ca = Some(keypair.cert.clone());
     config.keypair = Some(keypair);
 
-    info!("{:?}", config);
+    debug!(
+        "{}:{} started",
+        stablebabble(&addr.ip().octets()),
+        addr.port()
+    );
 
-    server::run(config).await.expect("server");
+    if matches.opt_present("c") {
+        client::run(config).await.expect("client");
+    } else {
+        server::run(config).await.expect("server");
+    }
 }

src/server.rs

@@ -13,10 +13,19 @@
 // OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 
 use crate::Config;
-use std::{io::Result};
+use std::{
-use tokio::{net::TcpListener, prelude::*};
+    io::{Error, ErrorKind, Result},
+    path::{PathBuf, MAIN_SEPARATOR},
+};
+use tokio::{
+    fs::{remove_file, File},
+    io::{self, AsyncReadExt},
+    net::TcpListener,
+};
 use tokio_libtls::prelude::*;
 
+const MAX_COPIES: usize = 1024;
+
 pub(crate) async fn run(config: Config) -> Result<()> {
     let (cert, key, ca) = config.load_keys()?;
     let mut options = config.load_server_options();
@@ -32,16 +41,146 @@ pub(crate) async fn run(config: Config) -> Result<()> {
 
     loop {
         let (tcp, _) = listener.accept().await?;
+        let size_limit = config.size_limit as u64;
+        let peer_addr: String = tcp.peer_addr()?.to_string();
         let options = options.build();
-        let mut tls = AsyncTls::accept_stream(tcp, &tls_config, options).await?;
+        let mut 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;
+            }
+        };
 
         tokio::spawn(async move {
-            loop {
+            let mut filename = PathBuf::from("/home/reyk/Downloads");
-                let mut buf = vec![0u8; 1024];
+
-                if !tls.read(&mut buf).await.is_ok() || !tls.write_all(&buf).await.is_ok() {
+            // Read and validate the filename
+            let line = match read_line(&peer_addr, &mut tls).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 tls).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;
+            }
+
+            debug!(
+                "{} status: receiving {} ({} bytes)",
+                peer_addr,
+                filename.display(),
+                file_size
+            );
+
+            // Create output file
+            let mut file = match File::create(&filename).await {
+                Ok(f) => f,
+                Err(err) => {
+                    debug!(
+                        "{} failed {}: file ({})",
+                        peer_addr,
+                        filename.display(),
+                        err
+                    );
+                    return;
+                }
+            };
+
+            // I/O
+            let mut reader = tls.take(file_size);
+            let copied = match io::copy(&mut reader, &mut file).await {
+                Ok(s) => s,
+                Err(err) => {
+                    debug!("{} failed: I/O ({})", peer_addr, err);
+                    return;
+                }
+            };
+
+            if copied != file_size {
+                drop(file);
+                let _ = remove_file("a.txt").await.is_ok();
+                warn!(
+                    "{} failed: {} ({}/{} bytes)",
+                    peer_addr,
+                    filename.display(),
+                    copied,
+                    file_size
+                );
+            } else {
+                info!(
+                    "{} success: {} ({} bytes)",
+                    peer_addr,
+                    filename.display(),
+                    copied
+                );
             }
         });
     }
 }
+
+async fn read_line<T: io::AsyncRead + Unpin>(peer: &str, reader: &mut T) -> Result<String> {
+    let mut buf = vec![0u8; 1024];
+    if let Err(err) = reader.read(&mut buf).await {
+        return Err(Error::new(
+            ErrorKind::Other,
+            format!("{} failed: read ({})", peer, err),
+        ));
+    }
+    let line = match String::from_utf8(buf) {
+        Ok(s) => s,
+        Err(err) => {
+            return Err(Error::new(
+                ErrorKind::Other,
+                format!("{} read failed: line ({})", peer, err),
+            ));
+        }
+    };
+    let len = line.find(|c: char| c == '\r' || c == '\n').unwrap_or(0);
+    Ok((&line[0..len]).to_owned())
+}