From da1b2cac2940eb1cf7f67328762fd5db5e1a0c8d Mon Sep 17 00:00:00 2001
From: Runxi Yu <me@runxiyu.org>
Date: Sun, 8 Dec 2024 10:04:03 +0800
Subject: Client teardown and proper nick changes

---
 clients.go  | 12 ++++++++++++
 cmd_nick.go | 16 +++++++++++++---
 main.go     |  1 +
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/clients.go b/clients.go
index a181f99..53fd134 100644
--- a/clients.go
+++ b/clients.go
@@ -45,6 +45,18 @@ func (client Client) ServerSource() string {
 	return client.UID
 }
 
+func (client *Client) Teardown() {
+	if client.conn != nil {
+		(*client.conn).Close()
+	}
+	if !uidToClient.CompareAndDelete(client.UID, client) {
+		slog.Error("uid inconsistent", "uid", client.UID, "client", client)
+	}
+	if !nickToClient.CompareAndDelete(client.Nick, client) {
+		slog.Error("nick inconsistent", "nick", client.Nick, "client", client)
+	}
+}
+
 type ClientState uint8
 
 const (
diff --git a/cmd_nick.go b/cmd_nick.go
index ad68cc6..89ee4d9 100644
--- a/cmd_nick.go
+++ b/cmd_nick.go
@@ -1,5 +1,9 @@
 package main
 
+import (
+	"log/slog"
+)
+
 func init() {
 	commandHandlers["NICK"] = handleClientNick
 }
@@ -9,14 +13,20 @@ func handleClientNick(msg RMsg, client *Client) bool {
 		client.Send(MakeMsg(self, ERR_NEEDMOREPARAMS, "NICK", "Not enough parameters"))
 		return true
 	}
-	_, exists := nickToClient.LoadOrStore(msg.Params[0], client)
+	already, exists := nickToClient.LoadOrStore(msg.Params[0], client)
 	if exists {
-		client.Send(MakeMsg(self, ERR_NICKNAMEINUSE, client.Nick, msg.Params[0], "Nickname is already in use"))
+		if already != client {
+			client.Send(MakeMsg(self, ERR_NICKNAMEINUSE, client.Nick, msg.Params[0], "Nickname is already in use"))
+		}
 	} else {
-		client.Nick = msg.Params[0]
 		if client.State == ClientStateRegistered {
+			if !nickToClient.CompareAndDelete(client.Nick, client) {
+				slog.Error("nick inconsistent", "nick", client.Nick, "client", client)
+				return false
+			}
 			client.Send(MakeMsg(client, "NICK", msg.Params[0]))
 		}
+		client.Nick = msg.Params[0]
 	}
 	return true
 }
diff --git a/main.go b/main.go
index 75afd08..b730b9d 100644
--- a/main.go
+++ b/main.go
@@ -40,6 +40,7 @@ func main() {
 		// TODO: Add to the UID table and make actually unique UIDs
 		go func() {
 			defer func() {
+				client.Teardown()
 				(*client.conn).Close()
 				// TODO: Unified client clean-up
 			}()
-- 
cgit v1.2.3