From 5ff2a77d5c00b7b516f1380e9008a484c806fe06 Mon Sep 17 00:00:00 2001
From: Trevor Arjeski <tmarjeski@gmail.com>
Date: Sat, 6 Jun 2020 23:14:22 +0300
Subject: [PATCH] Implemented DCC get for downloading files. Does not yet
 support sending files. Closes #206

Added connection timeout

Fixed NOTICE format which was causing an issue. Spawn regular task instead of local for file download.

Only read number of bytes read.
---
 README.md                   |   6 +-
 libtiny_client/src/dcc.rs   | 125 ++++++++++++++++++++++++++++
 libtiny_client/src/lib.rs   | 159 ++++++++++++++++++++++++++++++++++++
 libtiny_client/src/state.rs |  37 ++++++++-
 libtiny_wire/src/lib.rs     |   7 ++
 tiny/config.yml             |   1 +
 tiny/src/cmd.rs             |  62 +++++++++++++-
 tiny/src/config.rs          |   1 +
 tiny/src/conn.rs            |  33 +++++++-
 9 files changed, 427 insertions(+), 4 deletions(-)
 create mode 100644 libtiny_client/src/dcc.rs

diff --git a/README.md b/README.md
index a49c8c37..2a694515 100644
--- a/README.md
+++ b/README.md
@@ -195,7 +195,11 @@ Commands start with `/` character.
   Running this command in a server tab applies it to all channels of that
   server. You can check your notify state in the status line.
 
-## Server commands
+- `/dcc [get] <sender> <filename>`: When you receive a DCC SEND message, you can 
+  accept it and download the file with this command. Currently sending files
+  is not supported.
+
+## Server Commands
 
 For commands not supported by tiny as a slash command, sending the command in
 the server tab will send the message directly to the server.
diff --git a/libtiny_client/src/dcc.rs b/libtiny_client/src/dcc.rs
new file mode 100644
index 00000000..0200aa7b
--- /dev/null
+++ b/libtiny_client/src/dcc.rs
@@ -0,0 +1,125 @@
+use std::fmt;
+use std::net::{IpAddr, Ipv4Addr, SocketAddr};
+use std::str::FromStr;
+
+#[derive(Debug)]
+pub struct DCCRecord {
+    /// SEND or CHAT (only supporting SEND right now)
+    dcc_type: DCCType,
+    /// IP address and port
+    address: SocketAddr,
+    /// Nickname of person who wants to send
+    origin: String,
+    /// Client nickname
+    receiver: String,
+    /// Argument - filename or string "chat"
+    argument: String,
+    /// File size of file that will be sent in bytes
+    file_size: Option<u32>,
+}
+
+#[derive(Debug, Copy, Clone)]
+pub enum DCCType {
+    SEND,
+    CHAT,
+}
+
+#[derive(Debug)]
+pub struct DCCTypeParseError;
+
+impl std::error::Error for DCCTypeParseError {}
+
+impl fmt::Display for DCCTypeParseError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "DCCType could not be parsed from string.")
+    }
+}
+
+impl FromStr for DCCType {
+    type Err = DCCTypeParseError;
+    fn from_str(input: &str) -> Result<DCCType, DCCTypeParseError> {
+        match input {
+            "SEND" => Ok(DCCType::SEND),
+            "CHAT" => Ok(DCCType::CHAT),
+            _ => Err(DCCTypeParseError),
+        }
+    }
+}
+
+impl fmt::Display for DCCType {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            DCCType::CHAT => write!(f, "CHAT"),
+            DCCType::SEND => write!(f, "SEND"),
+        }
+    }
+}
+
+impl DCCRecord {
+    /// DCC type argument address port [size]
+    ///
+    /// type	 - The connection type.
+    /// argument - The connectin type dependant argument.
+    /// address	 - The host address of the initiator as an integer.
+    /// port	 - The port or the socket on which the initiator expects
+    /// 	       to receive the connection.
+    /// size     - If the connection type is "SEND" (see below), then size
+    /// 	       will indicate the size of the file being offered. Obsolete
+    /// 	       IRCII clients do not send this, so be prepared if this is
+    /// 	       not present.
+
+    /// The following DCC connection types are defined:
+    ///
+    /// Type	Purpose					                Argument
+    /// CHAT	To carry on a semi-secure conversation	the string "chat"
+    /// SEND	To send a file to the recipient		    the file name
+    pub fn from(
+        origin: &str,
+        receiver: &str,
+        msg: &str,
+    ) -> Result<DCCRecord, Box<dyn std::error::Error>> {
+        let mut param_iter = msg.split_whitespace();
+        let dcc_type: DCCType = param_iter.next().unwrap().parse()?;
+        let argument = param_iter.next().unwrap();
+        let address: u32 = param_iter.next().unwrap().parse()?;
+        let address_dot_decimal: Ipv4Addr = Ipv4Addr::new(
+            (address >> 24) as u8,
+            (address >> 16) as u8,
+            (address >> 8) as u8,
+            (address) as u8,
+        );
+        let port: u16 = param_iter.next().unwrap().parse()?;
+        let socket_address = SocketAddr::new(IpAddr::V4(address_dot_decimal), port);
+
+        let file_size = param_iter.next().and_then(|fs| fs.parse::<u32>().ok());
+
+        Ok(DCCRecord {
+            dcc_type,
+            address: socket_address,
+            origin: origin.to_string(),
+            receiver: receiver.to_string(),
+            argument: argument.to_string(),
+            file_size,
+        })
+    }
+
+    pub fn get_type(&self) -> DCCType {
+        self.dcc_type
+    }
+
+    pub fn get_addr(&self) -> &SocketAddr {
+        &self.address
+    }
+
+    pub fn get_argument(&self) -> &String {
+        &self.argument
+    }
+
+    pub fn get_file_size(&self) -> Option<u32> {
+        self.file_size
+    }
+
+    pub fn get_receiver(&self) -> &String {
+        &self.receiver
+    }
+}
diff --git a/libtiny_client/src/lib.rs b/libtiny_client/src/lib.rs
index 1f703b75..061ef95d 100644
--- a/libtiny_client/src/lib.rs
+++ b/libtiny_client/src/lib.rs
@@ -3,11 +3,13 @@
 #![allow(clippy::unneeded_field_pattern)]
 #![allow(clippy::cognitive_complexity)]
 
+mod dcc;
 mod pinger;
 mod state;
 mod stream;
 mod utils;
 
+pub use dcc::DCCType;
 pub use libtiny_wire as wire;
 
 use pinger::Pinger;
@@ -18,9 +20,12 @@ use futures::future::FutureExt;
 use futures::stream::{Fuse, StreamExt};
 use futures::{pin_mut, select};
 use std::net::{SocketAddr, ToSocketAddrs};
+use std::path::PathBuf;
 use std::time::Duration;
+use tokio::fs::OpenOptions;
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::sync::mpsc;
+use tokio::time::timeout;
 
 #[macro_use]
 extern crate log;
@@ -249,6 +254,36 @@ impl Client {
     pub fn get_chan_nicks(&self, chan: &str) -> Vec<String> {
         self.state.get_chan_nicks(chan)
     }
+
+    /// Checks if we have a stored DCCRecord for the provided parameters
+    /// Starts download of file if a record is found.
+    /// Deletes record after download attempt
+    /// Returns boolean if record is found
+    pub fn get_dcc_rec(&mut self, origin: &str, download_dir: PathBuf, file_name: &str) -> bool {
+        let rec = self.state.get_dcc_rec(origin, file_name);
+        if let Some(rec) = rec {
+            debug!("found rec {:?}", rec);
+            dcc_file_get(
+                &mut self.msg_chan,
+                *rec.get_addr(),
+                download_dir,
+                file_name.to_string(),
+                rec.get_receiver().to_string(),
+            );
+            true
+        } else {
+            false
+        }
+    }
+
+    /// Creates a dcc record and returns the argument (filename) and file size if provided
+    pub fn create_dcc_rec(
+        &self,
+        origin: &str,
+        msg: &str,
+    ) -> Option<(dcc::DCCType, String, Option<u32>)> {
+        self.state.add_dcc_rec(origin, msg)
+    }
 }
 
 //
@@ -660,3 +695,127 @@ async fn try_connect(
 
     None
 }
+
+/// Spawns a task to connect to given address to download
+/// a file
+/// https://www.irchelp.org/protocol/ctcpspec.html
+fn dcc_file_get(
+    msg_chan: &mut mpsc::Sender<Cmd>,
+    addr: SocketAddr,
+    mut download_dir: PathBuf,
+    file_name: String,
+    target: String,
+) {
+    let mut msg_chan = msg_chan.clone();
+    tokio::task::spawn(async move {
+        let stream = match timeout(
+            tokio::time::Duration::from_secs(25),
+            tokio::net::TcpStream::connect(addr),
+        )
+        .await
+        {
+            Err(err) => {
+                error!("{}", err);
+                let err_msg = "File transfer connection timed out.";
+                dcc_to_msgchan(&mut msg_chan, target.as_str(), err_msg).await;
+                return;
+            }
+            Ok(stream) => match stream {
+                Err(err) => {
+                    error!("{}", err);
+                    let err_msg = "File transfer connection failed.";
+                    dcc_to_msgchan(&mut msg_chan, target.as_str(), err_msg).await;
+                    return;
+                }
+                Ok(stream) => stream,
+            },
+        };
+
+        let (mut rx, mut tx) = tokio::io::split(stream);
+
+        // Try to create the download directory if it isn't there
+        if let Err(err) = tokio::fs::create_dir_all(&download_dir).await {
+            // Directory already exists
+            if err.kind() != tokio::io::ErrorKind::AlreadyExists {
+                error!("{}", err);
+                let err_msg = "Couldn't create download directory.";
+                dcc_to_msgchan(&mut msg_chan, target.as_str(), err_msg).await;
+                return;
+            }
+        }
+        // Add the file name to the file path
+        download_dir.push(&file_name);
+        // Create a new file for write inside the download directory
+        let mut file = match OpenOptions::new()
+            .write(true)
+            .create_new(true)
+            .open(download_dir)
+            .await
+        {
+            Err(err) => {
+                error!("{}", err);
+                let err_msg = format!("Couldn't create file: {}.", file_name);
+                dcc_to_msgchan(&mut msg_chan, target.as_str(), err_msg.as_str()).await;
+                return;
+            }
+            Ok(file) => file,
+        };
+        // Stores total bytes received
+        let mut bytes_received: u32 = 0;
+        loop {
+            let mut buf: [u8; 1024] = [0; 1024];
+            match rx.read(&mut buf).await {
+                Err(err) => {
+                    error!("{}", err);
+                    let err_msg = "Error reading from socket for dcc.";
+                    dcc_to_msgchan(&mut msg_chan, target.as_str(), err_msg).await;
+                    return;
+                }
+                Ok(0) => break,
+                Ok(bytes) => {
+                    bytes_received += bytes as u32;
+                    // Write bytes to file
+                    if let Err(err) = file.write(&buf[..bytes]).await {
+                        error!("{}", err);
+                        let err_msg = "Error writing to file.";
+                        dcc_to_msgchan(&mut msg_chan, target.as_str(), err_msg).await;
+                        return;
+                    }
+                    // Send back number of bytes read
+                    let response: [u8; 4] = [
+                        (bytes_received >> 24) as u8,
+                        (bytes_received >> 16) as u8,
+                        (bytes_received >> 8) as u8,
+                        (bytes_received) as u8,
+                    ];
+                    if let Err(err) = tx.write(&response).await {
+                        error!("{}", err);
+                        let err_msg = "Error responding to sender.";
+                        dcc_to_msgchan(&mut msg_chan, target.as_str(), err_msg).await;
+                        return;
+                    }
+                }
+            }
+        }
+        // sync file to filesystem
+        if let Err(err) = file.sync_all().await {
+            error!("{}", err);
+            let err_msg = "Error syncing file to disk.";
+            dcc_to_msgchan(&mut msg_chan, target.as_str(), err_msg).await;
+            return;
+        }
+        // Write message to client that file was successfully downloaded
+        let success_msg = format!(
+            "Downloaded file: {} ({1:.2}KB)",
+            file_name,
+            (bytes_received as f32 / 1024 as f32)
+        );
+        dcc_to_msgchan(&mut msg_chan, target.as_str(), success_msg.as_str()).await;
+    });
+}
+
+async fn dcc_to_msgchan(msg_chan: &mut mpsc::Sender<Cmd>, target: &str, msg: &str) {
+    msg_chan
+        .try_send(Cmd::Msg(wire::notice(target, msg)))
+        .unwrap();
+}
diff --git a/libtiny_client/src/state.rs b/libtiny_client/src/state.rs
index 68f34885..7f6d1f55 100644
--- a/libtiny_client/src/state.rs
+++ b/libtiny_client/src/state.rs
@@ -1,12 +1,13 @@
 #![allow(clippy::zero_prefixed_literal)]
 
+use crate::dcc::{DCCRecord, DCCType};
 use crate::utils;
 use crate::{Event, ServerInfo};
 use libtiny_wire as wire;
 use libtiny_wire::{Msg, Pfx};
 
 use std::cell::RefCell;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::rc::Rc;
 use tokio::sync::mpsc::Sender;
 
@@ -64,6 +65,36 @@ impl State {
     pub(crate) fn get_chan_nicks(&self, chan: &str) -> Vec<String> {
         self.inner.borrow().get_chan_nicks(chan)
     }
+
+    /// Gets a DCCRecord and removes it
+    pub(crate) fn get_dcc_rec(&self, origin: &str, argument: &str) -> Option<DCCRecord> {
+        self.inner
+            .borrow_mut()
+            .dcc_recs
+            .remove(&(origin.to_string(), argument.to_string()))
+    }
+
+    /// Adds a DCC record and returns the type, argument and filesize if successfully created
+    pub(crate) fn add_dcc_rec(
+        &self,
+        origin: &str,
+        msg: &str,
+    ) -> Option<(DCCType, String, Option<u32>)> {
+        let record = DCCRecord::from(origin, &self.get_nick(), msg);
+        debug!("record: {:?}", record);
+        if let Ok(record) = record {
+            let argument = record.get_argument().clone();
+            let dcc_type = record.get_type().clone();
+            let file_size = record.get_file_size();
+            self.inner
+                .borrow_mut()
+                .dcc_recs
+                .insert((origin.to_string(), argument.to_string()), record);
+            Some((dcc_type, argument, file_size))
+        } else {
+            None
+        }
+    }
 }
 
 struct StateInner {
@@ -91,6 +122,9 @@ struct StateInner {
     /// order, in TUI?
     chans: Vec<(String, HashSet<String>)>,
 
+    /// Map of DCCRecords indexed by the nickname of the sender and the argument (filename or 'chat')
+    dcc_recs: HashMap<(String, String), DCCRecord>,
+
     /// Away reason if away mode is on. `None` otherwise. TODO: I don't think the message is used?
     away_status: Option<String>,
 
@@ -124,6 +158,7 @@ impl StateInner {
             current_nick_idx: 0,
             current_nick,
             chans,
+            dcc_recs: HashMap::new(),
             away_status: None,
             servername: None,
             usermask: None,
diff --git a/libtiny_wire/src/lib.rs b/libtiny_wire/src/lib.rs
index c3a24e5b..06287a81 100644
--- a/libtiny_wire/src/lib.rs
+++ b/libtiny_wire/src/lib.rs
@@ -49,6 +49,11 @@ pub fn privmsg(msgtarget: &str, msg: &str) -> String {
     format!("PRIVMSG {} :{}\r\n", msgtarget, msg)
 }
 
+pub fn notice(msgtarget: &str, msg: &str) -> String {
+    assert!(msgtarget.len() + msg.len() + 12 <= 512);
+    format!("NOTICE {} :{}\r\n", msgtarget, msg)
+}
+
 pub fn action(msgtarget: &str, msg: &str) -> String {
     assert!(msgtarget.len() + msg.len() + 21 <= 512); // See comments in `privmsg`
     format!("PRIVMSG {} :\x01ACTION {}\x01\r\n", msgtarget, msg)
@@ -116,6 +121,7 @@ pub struct Msg {
 pub enum CTCP {
     Version,
     Action,
+    DCC,
     Other(String),
 }
 
@@ -124,6 +130,7 @@ impl CTCP {
         match s {
             "VERSION" => CTCP::Version,
             "ACTION" => CTCP::Action,
+            "DCC" => CTCP::DCC,
             _ => CTCP::Other(s.to_owned()),
         }
     }
diff --git a/tiny/config.yml b/tiny/config.yml
index 30553170..a519b13e 100644
--- a/tiny/config.yml
+++ b/tiny/config.yml
@@ -35,6 +35,7 @@ defaults:
     realname: yourname
     join: []
     tls: false
+    download_dir: '/tmp/tiny_downloads'
 
 # Where to put log files
 log_dir: '{}'
diff --git a/tiny/src/cmd.rs b/tiny/src/cmd.rs
index 7b87bf02..37e727a7 100644
--- a/tiny/src/cmd.rs
+++ b/tiny/src/cmd.rs
@@ -97,7 +97,7 @@ fn find_client<'a>(clients: &'a mut Vec<Client>, serv_name: &str) -> Option<&'a
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
-static CMDS: [&Cmd; 9] = [
+static CMDS: [&Cmd; 10] = [
     &AWAY_CMD,
     &CLOSE_CMD,
     &CONNECT_CMD,
@@ -106,6 +106,7 @@ static CMDS: [&Cmd; 9] = [
     &MSG_CMD,
     &NAMES_CMD,
     &NICK_CMD,
+    &DCC_CMD,
     &HELP_CMD,
 ];
 
@@ -491,6 +492,65 @@ fn help(args: CmdArgs) {
     }
 }
 
+static DCC_CMD: Cmd = Cmd {
+    name: "dcc",
+    cmd_fn: dcc,
+    description: "Accepts DCC request",
+    usage: "/dcc get <sender> <filename>"
+};
+
+const GET: &str = "get";
+const SEND: &str = "send";
+
+fn dcc(args: CmdArgs) {
+    let CmdArgs {
+        args,
+        ui,
+        clients,
+        src,
+        defaults,
+        ..
+    } = args;
+
+    if let Some(client) = find_client(clients, &src.serv_name()) {
+        // parse args, ex. GET <sender nick> <filename>
+        let mut params = args.split_whitespace();
+        if let Some(sub_cmd) = params.next() {
+            match sub_cmd {
+                GET => {
+                    if let Some(origin) = params.next() {
+                        if let Some(file_name) = params.next() {
+                            let found = client.get_dcc_rec(
+                                origin,
+                                defaults.download_dir.clone(),
+                                file_name,
+                            );
+                            match found {
+                                true => ui.add_client_msg(
+                                    "DCC Record found...starting download",
+                                    &MsgTarget::CurrentTab,
+                                ),
+                                false => ui.add_client_msg(
+                                    "DCC Record not found.",
+                                    &MsgTarget::CurrentTab,
+                                ),
+                            }
+                        }
+                    }
+                }
+                // TODO: Support this if people request it
+                SEND => ui.add_client_err_msg(
+                    "SEND command not implemented yet.",
+                    &MsgTarget::CurrentTab,
+                ),
+                _ => ui.add_client_err_msg(
+                    DCC_CMD.usage,
+                    &MsgTarget::CurrentTab,
+                ),
+            }
+        }
+    }
+}
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 #[cfg(test)]
diff --git a/tiny/src/config.rs b/tiny/src/config.rs
index ac64d1b9..553d3781 100644
--- a/tiny/src/config.rs
+++ b/tiny/src/config.rs
@@ -58,6 +58,7 @@ pub(crate) struct Defaults {
     pub(crate) join: Vec<String>,
     #[serde(default)]
     pub(crate) tls: bool,
+    pub(crate) download_dir: PathBuf,
 }
 
 #[derive(Deserialize)]
diff --git a/tiny/src/conn.rs b/tiny/src/conn.rs
index 519bfea5..a421f31a 100644
--- a/tiny/src/conn.rs
+++ b/tiny/src/conn.rs
@@ -4,7 +4,7 @@
 //! IRC event handling
 
 use futures::stream::StreamExt;
-use libtiny_client::Client;
+use libtiny_client::{Client, DCCType};
 use libtiny_ui::{MsgTarget, TabStyle, UI};
 use libtiny_wire as wire;
 use tokio::sync::mpsc;
@@ -155,6 +155,37 @@ fn handle_irc_msg(ui: &dyn UI, client: &Client, msg: wire::Msg) {
                 return;
             }
 
+            if ctcp == Some(wire::CTCP::DCC) {
+                let msg_target = if ui.user_tab_exists(serv, origin) {
+                    MsgTarget::User { serv, nick: origin }
+                } else {
+                    MsgTarget::Server { serv }
+                };
+                if let Some((dcc_type, argument, file_size)) =
+                    client.create_dcc_rec(&origin.to_owned(), &msg)
+                {
+                    match dcc_type {
+                        DCCType::SEND => {
+                            let file_size =
+                                file_size.map_or("unknown size".to_string(), |s| s.to_string());
+                            ui.add_client_msg(
+                                &format!(
+                                    "{} wants to send file - {} ({} bytes)",
+                                    origin, argument, file_size
+                                ),
+                                &msg_target,
+                            );
+                            ui.add_client_msg(
+                                &format!("Command to accept: /dcc get {} {}", origin, argument),
+                                &msg_target,
+                            );
+                        }
+                        DCCType::CHAT => {}
+                    }
+                }
+                return;
+            }
+
             let is_action = ctcp == Some(wire::CTCP::Action);
 
             match target {