From a8eabe8e9c3e7616ad5a087168774eefe5c50b79 Mon Sep 17 00:00:00 2001
From: Daniel Oaks <daniel@danieloaks.net>
Date: Tue, 29 Nov 2016 22:33:10 +1000
Subject: [PATCH] client: Fix a lot of bugs around setting NICK

---
 irc/client.go            |  9 ++++++---
 irc/client_lookup_set.go | 26 +++++++++++++-------------
 irc/server.go            |  8 --------
 3 files changed, 19 insertions(+), 24 deletions(-)

diff --git a/irc/client.go b/irc/client.go
index 53dc23a5..fb738b6e 100644
--- a/irc/client.go
+++ b/irc/client.go
@@ -358,9 +358,12 @@ func (client *Client) SetNickname(nickname string) error {
 		return ErrNickAlreadySet
 	}
 
-	client.nick = nickname
-	client.updateNick()
-	return client.server.clients.Add(client)
+	err := client.server.clients.Add(client, nickname)
+	if err == nil {
+		client.nick = nickname
+		client.updateNick()
+	}
+	return err
 }
 
 // ChangeNickname changes the existing nickname of the client.
diff --git a/irc/client_lookup_set.go b/irc/client_lookup_set.go
index cdd4e350..57f01b58 100644
--- a/irc/client_lookup_set.go
+++ b/irc/client_lookup_set.go
@@ -48,8 +48,8 @@ func NewClientLookupSet() *ClientLookupSet {
 
 func (clients *ClientLookupSet) Count() int {
 	clients.ByNickMutex.Lock()
+	defer clients.ByNickMutex.Unlock()
 	count := len(clients.ByNick)
-	clients.ByNickMutex.Unlock()
 	return count
 }
 
@@ -60,8 +60,8 @@ func (clients *ClientLookupSet) Has(nick string) bool {
 		return false
 	}
 	clients.ByNickMutex.Lock()
+	defer clients.ByNickMutex.Unlock()
 	_, exists := clients.ByNick[casefoldedName]
-	clients.ByNickMutex.Unlock()
 	return exists
 }
 
@@ -79,23 +79,24 @@ func (clients *ClientLookupSet) Get(nick string) *Client {
 	casefoldedName, err := CasefoldName(nick)
 	if err == nil {
 		clients.ByNickMutex.Lock()
+		defer clients.ByNickMutex.Unlock()
 		cli := clients.ByNick[casefoldedName]
-		clients.ByNickMutex.Unlock()
 		return cli
 	}
 	return nil
 }
 
-func (clients *ClientLookupSet) Add(client *Client) error {
-	if !client.HasNick() {
-		return ErrNickMissing
+func (clients *ClientLookupSet) Add(client *Client, nick string) error {
+	nick, err := CasefoldName(nick)
+	if err != nil {
+		return err
 	}
 	clients.ByNickMutex.Lock()
 	defer clients.ByNickMutex.Unlock()
-	if clients.getNoMutex(client.nick) != nil {
+	if clients.getNoMutex(nick) != nil {
 		return ErrNicknameInUse
 	}
-	clients.ByNick[client.nickCasefolded] = client
+	clients.ByNick[nick] = client
 	return nil
 }
 
@@ -104,12 +105,11 @@ func (clients *ClientLookupSet) Remove(client *Client) error {
 		return ErrNickMissing
 	}
 	clients.ByNickMutex.Lock()
+	defer clients.ByNickMutex.Unlock()
 	if clients.getNoMutex(client.nick) != client {
-		clients.ByNickMutex.Unlock()
 		return ErrNicknameMismatch
 	}
 	delete(clients.ByNick, client.nickCasefolded)
-	clients.ByNickMutex.Unlock()
 	return nil
 }
 
@@ -150,6 +150,7 @@ func (clients *ClientLookupSet) AllWithCaps(caps ...Capability) (set ClientSet)
 	set = make(ClientSet)
 
 	clients.ByNickMutex.Lock()
+	defer clients.ByNickMutex.Unlock()
 	var client *Client
 	for _, client = range clients.ByNick {
 		// make sure they have all the required caps
@@ -161,7 +162,6 @@ func (clients *ClientLookupSet) AllWithCaps(caps ...Capability) (set ClientSet)
 
 		set.Add(client)
 	}
-	clients.ByNickMutex.Unlock()
 
 	return set
 }
@@ -176,12 +176,12 @@ func (clients *ClientLookupSet) FindAll(userhost string) (set ClientSet) {
 	matcher := ircmatch.MakeMatch(userhost)
 
 	clients.ByNickMutex.Lock()
+	defer clients.ByNickMutex.Unlock()
 	for _, client := range clients.ByNick {
 		if matcher.Match(client.nickMaskCasefolded) {
 			set.Add(client)
 		}
 	}
-	clients.ByNickMutex.Unlock()
 
 	return set
 }
@@ -195,13 +195,13 @@ func (clients *ClientLookupSet) Find(userhost string) *Client {
 	var matchedClient *Client
 
 	clients.ByNickMutex.Lock()
+	defer clients.ByNickMutex.Unlock()
 	for _, client := range clients.ByNick {
 		if matcher.Match(client.nickMaskCasefolded) {
 			matchedClient = client
 			break
 		}
 	}
-	clients.ByNickMutex.Unlock()
 
 	return matchedClient
 }
diff --git a/irc/server.go b/irc/server.go
index 8754db19..a32a9760 100644
--- a/irc/server.go
+++ b/irc/server.go
@@ -658,13 +658,6 @@ func userHandler(server *Server, client *Client, msg ircmsg.IrcMessage) bool {
 		return true
 	}
 
-	// set user info and log client in
-	//TODO(dan): Could there be a race condition here with adding/removing the client?
-	//TODO(dan): we should do something like server.clients.Replace(client) instead
-
-	// we do it this way to ONLY replace what hasn't already been set
-	server.clients.Remove(client)
-
 	if !client.HasUsername() {
 		client.username = "~" + msg.Params[0]
 		// don't bother updating nickmask here, it's not valid anyway
@@ -673,7 +666,6 @@ func userHandler(server *Server, client *Client, msg ircmsg.IrcMessage) bool {
 		client.realname = msg.Params[3]
 	}
 
-	server.clients.Add(client)
 	server.tryRegister(client)
 
 	return false