From d856c39e22ed1afdfe7c7d58d72eb318bf554f2b Mon Sep 17 00:00:00 2001
From: Jordan Doyle <jordan@doyle.la>
Date: Thu, 2 Feb 2023 00:48:06 +0000
Subject: [PATCH] Clean shutdowns on negotiation failures

---
 src/connection.rs | 24 +++++++++++++++++++++++-
 src/main.rs       | 28 ++++++++++++++++++++++++----
 2 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/src/connection.rs b/src/connection.rs
index ff37ef8..b1d707c 100644
--- a/src/connection.rs
+++ b/src/connection.rs
@@ -9,7 +9,9 @@ use std::{
 use actix::{io::FramedWrite, Actor, Addr};
 use const_format::concatcp;
 use futures::{SinkExt, TryStreamExt};
-use irc_proto::{error::ProtocolError, CapSubCommand, Command, IrcCodec, Message, Prefix};
+use irc_proto::{
+    error::ProtocolError, CapSubCommand, Command, IrcCodec, Message, Prefix, Response,
+};
 use tokio::{
     io::{ReadHalf, WriteHalf},
     net::TcpStream,
@@ -198,6 +200,10 @@ pub async fn negotiate_client_connection(
         .map_err(|e| ProtocolError::Io(Error::new(ErrorKind::InvalidData, e)))?;
 
     if !reserved_nick {
+        write
+            .send(NickNotOwnedByUser(initiated.nick).into_message())
+            .await?;
+
         return Err(ProtocolError::Io(Error::new(
             ErrorKind::InvalidData,
             "nick is already in use by another user",
@@ -207,6 +213,22 @@ pub async fn negotiate_client_connection(
     Ok(Some(initiated))
 }
 
+pub struct NickNotOwnedByUser(String);
+
+impl NickNotOwnedByUser {
+    #[must_use]
+    pub fn into_message(self) -> Message {
+        Message {
+            tags: None,
+            prefix: None,
+            command: Command::Response(
+                Response::ERR_NICKLOCKED,
+                vec![self.0, "You must use a nick assigned to you".to_string()],
+            ),
+        }
+    }
+}
+
 /// Return an acknowledgement to the client for their requested capabilities.
 pub struct AcknowledgedCapabilities(String);
 
diff --git a/src/main.rs b/src/main.rs
index bf5e2c0..2ce8d41 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -11,7 +11,8 @@ use actix::{io::FramedWrite, Actor, Addr, AsyncContext, Supervisor};
 use actix_rt::{Arbiter, System};
 use bytes::BytesMut;
 use clap::Parser;
-use irc_proto::IrcCodec;
+use futures::SinkExt;
+use irc_proto::{Command, IrcCodec, Message};
 use rand::seq::SliceRandom;
 use sqlx::migrate::Migrator;
 use tokio::{
@@ -147,9 +148,28 @@ async fn start_tcp_acceptor_loop(
 
             // ensure we have all the details required to actually connect the client to the server
             // (ie. we have a nick, user, etc)
-            let Some(connection) = connection::negotiate_client_connection(&mut read, &mut write, addr, &persistence, database).await.unwrap() else {
-                error!("Failed to fully handshake with client, dropping connection");
-                return;
+            let connection = match connection::negotiate_client_connection(&mut read, &mut write, addr, &persistence, database).await {
+                Ok(Some(v)) => v,
+                Ok(None) => {
+                    error!("Failed to fully handshake with client, dropping connection");
+
+                    let command = Command::ERROR("You must use SASL to connect to this server".to_string());
+                    if let Err(error) = write.send(Message { tags: None, prefix: None, command, }).await {
+                        error!(%error, "Failed to send error message to client, forcefully closing connection.");
+                    }
+
+                    return;
+                }
+                Err(error) => {
+                    error!(%error, "An error occurred whilst handshaking with client");
+
+                    let command = Command::ERROR(error.to_string());
+                    if let Err(error) = write.send(Message { tags: None, prefix: None, command, }).await {
+                        error!(%error, "Failed to send error message to client, forcefully closing connection.");
+                    }
+
+                    return;
+                }
             };
 
             // spawn the client's actor
--
libgit2 1.7.2