2 files changed, 52 insertions, 34 deletions

diff --git a/src/modules/m_spanningtree/treesocket.h b/src/modules/m_spanningtree/treesocket.h
index a06ea0019..b775905c0 100644
--- a/src/modules/m_spanningtree/treesocket.h
+++ b/src/modules/m_spanningtree/treesocket.h
@@ -117,6 +117,18 @@ class TreeSocket : public BufferedSocket
 	/** Send all additional info about the given server to this server */
 	void SendServerInfo(TreeServer* from);
 
+	/** Find the User source of a command given a prefix and a command string.
+	 * This connection must be fully up when calling this function.
+	 * @param prefix Prefix string to find the source User object for. Can be a sid, a uuid or a server name.
+	 * @param command The command whose source to find. This is required because certain commands (like mode
+	 * changes and kills) must be processed even if their claimed source doesn't exist. If the given command is
+	 * such a command and the source does not exist, the function returns a valid FakeUser that can be used to
+	 * to process the command with.
+	 * @return The command source to use when processing the command or NULL if the source wasn't found.
+	 * Note that the direction of the returned source is not verified.
+	 */
+	User* FindSource(const std::string& prefix, const std::string& command);
+
  public:
 	const time_t age;
 
diff --git a/src/modules/m_spanningtree/treesocket2.cpp b/src/modules/m_spanningtree/treesocket2.cpp
index d16e55702..3df0bdee8 100644
--- a/src/modules/m_spanningtree/treesocket2.cpp
+++ b/src/modules/m_spanningtree/treesocket2.cpp
@@ -226,46 +226,52 @@ void TreeSocket::ProcessLine(std::string &line)
 	}
 }
 
-void TreeSocket::ProcessConnectedLine(std::string& prefix, std::string& command, parameterlist& params)
+User* TreeSocket::FindSource(const std::string& prefix, const std::string& command)
 {
+	// Empty prefix means the source is the directly connected server that sent this command
+	if (prefix.empty())
+		return MyRoot->ServerUser;
+
+	// If the prefix string is a uuid or a sid FindUUID() returns the appropriate User object
 	User* who = ServerInstance->FindUUID(prefix);
+	if (who)
+		return who;
 
-	if (!who)
-	{
-		TreeServer* ServerSource = Utils->FindServer(prefix);
-		if (prefix.empty())
-			ServerSource = MyRoot;
+	// Some implementations wrongly send a server name as prefix occasionally, handle that too for now
+	TreeServer* const server = Utils->FindServer(prefix);
+	if (server)
+		return server->ServerUser;
 
-		if (ServerSource)
-		{
-			who = ServerSource->ServerUser;
-		}
-		else
-		{
-			/* It is important that we don't close the link here, unknown prefix can occur
-			 * due to various race conditions such as the KILL message for a user somehow
-			 * crossing the users QUIT further upstream from the server. Thanks jilles!
-			 */
+	/* It is important that we don't close the link here, unknown prefix can occur
+	 * due to various race conditions such as the KILL message for a user somehow
+	 * crossing the users QUIT further upstream from the server. Thanks jilles!
+	 */
 
-			if ((prefix.length() == UIDGenerator::UUID_LENGTH) && (isdigit(prefix[0])) &&
-				((command == "FMODE") || (command == "MODE") || (command == "KICK") || (command == "TOPIC") || (command == "KILL") || (command == "ADDLINE") || (command == "DELLINE")))
-			{
-				/* Special case, we cannot drop these commands as they've been committed already on a
-				 * part of the network by the time we receive them, so in this scenario pretend the
-				 * command came from a server to avoid desync.
-				 */
+	if ((prefix.length() == UIDGenerator::UUID_LENGTH) && (isdigit(prefix[0])) &&
+		((command == "FMODE") || (command == "MODE") || (command == "KICK") || (command == "TOPIC") || (command == "KILL") || (command == "ADDLINE") || (command == "DELLINE")))
+	{
+		/* Special case, we cannot drop these commands as they've been committed already on a
+		 * part of the network by the time we receive them, so in this scenario pretend the
+		 * command came from a server to avoid desync.
+		 */
+
+		who = ServerInstance->FindUUID(prefix.substr(0, 3));
+		if (who)
+			return who;
+		return this->MyRoot->ServerUser;
+	}
 
-				who = ServerInstance->FindUUID(prefix.substr(0, 3));
-				if (!who)
-					who = this->MyRoot->ServerUser;
-			}
-			else
-			{
-				ServerInstance->Logs->Log(MODNAME, LOG_DEBUG, "Command '%s' from unknown prefix '%s'! Dropping entire command.",
-					command.c_str(), prefix.c_str());
-				return;
-			}
-		}
+	// Unknown prefix
+	return NULL;
+}
+
+void TreeSocket::ProcessConnectedLine(std::string& prefix, std::string& command, parameterlist& params)
+{
+	User* who = FindSource(prefix, command);
+	if (!who)
+	{
+		ServerInstance->Logs->Log(MODNAME, LOG_DEBUG, "Command '%s' from unknown prefix '%s'! Dropping entire command.", command.c_str(), prefix.c_str());
+		return;
 	}
 
 	/*