From 87353f7df9c55c2871f0b29186769d531c525e7d Mon Sep 17 00:00:00 2001
From: Wiebe Cazemier <wiebe@halfgaar.net>
Date: Mon, 16 Oct 2023 10:07:57 +0200
Subject: [PATCH] Fix denying wills on authentication failure

---
 client.cpp     | 14 ++++++++++----
 client.h       |  5 ++++-
 mqttpacket.cpp |  5 +----
 threaddata.cpp |  7 +++++++
 4 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/client.cpp b/client.cpp
index a094075e..f96e41ca 100644
--- a/client.cpp
+++ b/client.cpp
@@ -883,11 +883,16 @@ void Client::setClientProperties(bool connectPacketSeen, uint16_t keepalive, uin
     this->supportsRetained = supportsRetained;
 }
 
-void Client::setWill(WillPublish &&willPublish)
+void Client::stageWill(WillPublish &&willPublish)
 {
-    this->willPublish = std::make_shared<WillPublish>(std::move(willPublish));
-    this->willPublish->client_id = this->clientid;
-    this->willPublish->username = this->username;
+    this->stagedWillPublish = std::make_shared<WillPublish>(std::move(willPublish));
+    this->stagedWillPublish->client_id = this->clientid;
+    this->stagedWillPublish->username = this->username;
+}
+
+void Client::setWillFromStaged()
+{
+    this->willPublish = std::move(stagedWillPublish);
 }
 
 void Client::assignSession(std::shared_ptr<Session> &session)
@@ -942,6 +947,7 @@ std::chrono::seconds Client::getSecondsTillKeepAliveAction() const
 void Client::clearWill()
 {
     willPublish.reset();
+    stagedWillPublish.reset();
 
     if (session)
         session->clearWill();
diff --git a/client.h b/client.h
index cf107ecf..204843aa 100644
--- a/client.h
+++ b/client.h
@@ -86,6 +86,7 @@ class Client
     uint16_t keepalive = 10;
     bool clean_start = false;
 
+    std::shared_ptr<WillPublish> stagedWillPublish;
     std::shared_ptr<WillPublish> willPublish;
 
     const int epoll_fd;
@@ -139,7 +140,8 @@ class Client
                              uint32_t maxOutgoingPacketSize, uint16_t maxOutgoingTopicAliasValue);
     void setClientProperties(bool connectPacketSeen, uint16_t keepalive, uint32_t maxOutgoingPacketSize, uint16_t maxOutgoingTopicAliasValue, bool supportsRetained);
     void setWill(const std::string &topic, const std::string &payload, bool retain, uint8_t qos);
-    void setWill(WillPublish &&willPublish);
+    void stageWill(WillPublish &&willPublish);
+    void setWillFromStaged();
     void clearWill();
     void setAuthenticated(bool value) { authenticated = value;}
     bool getAuthenticated() { return authenticated; }
@@ -148,6 +150,7 @@ class Client
     const std::string &getUsername() const { return this->username; }
     std::string &getMutableUsername();
     std::shared_ptr<WillPublish> &getWill() { return this->willPublish; }
+    const std::shared_ptr<WillPublish> &getStagedWill() { return this->stagedWillPublish; }
     void assignSession(std::shared_ptr<Session> &session);
     std::shared_ptr<Session> getSession();
     void setDisconnectReason(const std::string &reason);
diff --git a/mqttpacket.cpp b/mqttpacket.cpp
index 732cc289..bce746fc 100644
--- a/mqttpacket.cpp
+++ b/mqttpacket.cpp
@@ -875,10 +875,7 @@ void MqttPacket::handleConnect()
 
     if (settings.willsEnabled && connectData.will_flag)
     {
-        if (authentication.aclCheck(connectData.willpublish, connectData.willpublish.payload, AclAccess::register_will) == AuthResult::success)
-        {
-            sender->setWill(std::move(connectData.willpublish));
-        }
+        sender->stageWill(std::move(connectData.willpublish));
     }
 
     // Stage connack, for immediate or delayed use when auth succeeds.
diff --git a/threaddata.cpp b/threaddata.cpp
index 97b388ca..f2aa849b 100644
--- a/threaddata.cpp
+++ b/threaddata.cpp
@@ -182,6 +182,13 @@ void ThreadData::continuationOfAuthentication(std::shared_ptr<Client> &client, A
             if (!returnData.empty())
                 client->addAuthReturnDataToStagedConnAck(returnData);
 
+            const std::shared_ptr<WillPublish> will = client->getStagedWill();
+
+            if (will && authentication.aclCheck(*will, will->payload, AclAccess::register_will) == AuthResult::success)
+            {
+                client->setWillFromStaged();
+            }
+
             client->sendConnackSuccess();
             subscriptionStore->registerClientAndKickExistingOne(client);
         }