From aae97bda3ea3452e8a6a3da82d3cfdf0adb8e038 Mon Sep 17 00:00:00 2001
From: Sadie Powell <sadie@witchery.services>
Date: Tue, 2 May 2023 13:33:31 +0100
Subject: [PATCH] Add <cloak:class> to limit cloaks to a specific connect
 class.

---
 docs/conf/modules.conf.example |  8 ++++++++
 include/modules/cloak.h        | 17 ++++++++++++++++-
 src/modules/m_cloak_md5.cpp    |  4 ++--
 src/modules/m_cloak_sha256.cpp |  4 ++--
 4 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/docs/conf/modules.conf.example b/docs/conf/modules.conf.example
index 74643a53c..5aecd58de 100644
--- a/docs/conf/modules.conf.example
+++ b/docs/conf/modules.conf.example
@@ -561,12 +561,14 @@
 #
 #<cloak method="half"
 #       key="changeme"
+#       class=""
 #       domainparts="3"
 #       prefix="net-"
 #       ignorecase="no">
 #
 #<cloak method="full"
 #       key="changeme"
+#       class=""
 #       prefix="net-"
 #       ignorecase="no">
 
@@ -586,6 +588,10 @@
 #  key       -  The secret key to use when hashing hostnames. This    #
 #               MUST be at least 30 characters long.                  #
 #                                                                     #
+#  class     - If non-empty then a comma-delimited list of connect    #
+#              class names that a user has to be in to get the cloak  #
+#              from this tag.                                         #
+#                                                                     #
 #  prefix    - A freeform value to prefix cloaks with. This must not  #
 #              contain spaces.                                        #
 #                                                                     #
@@ -615,6 +621,7 @@
 #
 #<cloak method="hmac-sha256"
 #       key="changeme"
+#       class=""
 #       prefix="MyNet"
 #       suffix="ip"
 #       case="lower"
@@ -624,6 +631,7 @@
 #
 #<cloak method="hmac-sha256-ip"
 #       key="changeme"
+#       class=""
 #       prefix="MyNet"
 #       suffix="ip"
 #       case="lower"
diff --git a/include/modules/cloak.h b/include/modules/cloak.h
index 2d3f41fd6..7a2d7d525 100644
--- a/include/modules/cloak.h
+++ b/include/modules/cloak.h
@@ -117,10 +117,25 @@ private:
 	/** The name of the engine that created this method. */
 	std::string provname;
 
+	/** The connect classes that a user can be in before */
+	insp::flat_set<std::string> classes;
+
 protected:
-	Method(const Engine* engine) ATTR_NOT_NULL(2)
+	Method(const Engine* engine, const std::shared_ptr<ConfigTag>& tag) ATTR_NOT_NULL(2)
 		: provname(engine->name)
 	{
+		irc::commasepstream klassstream(tag->getString("class"));
+		for (std::string klass; klassstream.GetToken(klass); )
+			classes.insert(klass);
+	}
+
+	bool MatchesUser(LocalUser* user) const
+	{
+		if (!classes.empty() && !stdalgo::isin(classes, user->GetClass()->GetName()))
+			return false;
+
+		// All fields matched.
+		return true;
 	}
 
 public:
diff --git a/src/modules/m_cloak_md5.cpp b/src/modules/m_cloak_md5.cpp
index abbdfa4db..f09a651b4 100644
--- a/src/modules/m_cloak_md5.cpp
+++ b/src/modules/m_cloak_md5.cpp
@@ -75,7 +75,7 @@ struct CloakInfo final
 	std::string suffix;
 
 	CloakInfo(const Cloak::Engine* engine, const std::shared_ptr<ConfigTag>& tag, CloakMode Mode, const std::string& Key)
-		: Cloak::Method(engine)
+		: Cloak::Method(engine, tag)
 		, mode(Mode)
 		, domainparts(tag->getNum<unsigned int>("domainparts", 3, 1, 10))
 		, ignorecase(tag->getBool("ignorecase"))
@@ -253,7 +253,7 @@ struct CloakInfo final
 
 	std::string Generate(LocalUser* user) override ATTR_NOT_NULL(2)
 	{
-		if (!md5 || !user->client_sa.is_ip())
+		if (!md5 || !user->client_sa.is_ip() || !MatchesUser(user))
 			return {};
 
 		return GenCloak(user->client_sa, user->GetAddress(), user->GetRealHost());
diff --git a/src/modules/m_cloak_sha256.cpp b/src/modules/m_cloak_sha256.cpp
index 97038a569..a269048da 100644
--- a/src/modules/m_cloak_sha256.cpp
+++ b/src/modules/m_cloak_sha256.cpp
@@ -186,7 +186,7 @@ private:
 
 public:
 	SHA256Method(const Cloak::Engine* engine, const std::shared_ptr<ConfigTag>& tag, const std::string& k, psl_ctx_t* p, bool ch) ATTR_NOT_NULL(2)
-		: Cloak::Method(engine)
+		: Cloak::Method(engine, tag)
 		, cloakhost(ch)
 		, hostparts(tag->getNum<unsigned long>("hostparts", 3, 0, ServerInstance->Config->Limits.MaxHost / 2))
 		, key(k)
@@ -214,7 +214,7 @@ public:
 
 	std::string Generate(LocalUser* user) override ATTR_NOT_NULL(2)
 	{
-		if (!sha256)
+		if (!sha256 || !MatchesUser(user))
 			return {};
 
 		irc::sockets::sockaddrs sa(false);