From e367fb22642d5b0cd379b0605bb62127efa7909a Mon Sep 17 00:00:00 2001
From: Alfian Losari <alfianlosari@gmail.com>
Date: Tue, 20 Feb 2024 14:11:25 +0700
Subject: [PATCH] Handle courier event message send success event not being
 triggered for QoS 3/4 (#31)

* Handle courier event message send success event not being triggered for qos 3 and 4
---
 .../ConnectionObservableObject.swift          |  2 +-
 CourierE2EApp/Views/ConnectionView.swift      |  2 +-
 .../MQTTClientFrameworkSessionManager.swift   | 22 ++++++++++++++++++-
 .../MQTTClientGJ/MQTTClientGJ/MQTTSession.m   | 13 ++++++++++-
 4 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/CourierE2EApp/View Models/ConnectionObservableObject.swift b/CourierE2EApp/View Models/ConnectionObservableObject.swift
index 535d058..bd46ba3 100644
--- a/CourierE2EApp/View Models/ConnectionObservableObject.swift	
+++ b/CourierE2EApp/View Models/ConnectionObservableObject.swift	
@@ -163,7 +163,7 @@ final class ConnectionObservableObject: ObservableObject {
 extension ConnectionObservableObject: ICourierEventHandler {
     
     func onEvent(_ event: CourierEvent) {
-//        print("EVENT: \(event.type)")
+        print("EVENT: \(event.type)")
         switch event.type {
         case .connectionSuccess:
             if connectionServiceProvider.isCleanSession {
diff --git a/CourierE2EApp/Views/ConnectionView.swift b/CourierE2EApp/Views/ConnectionView.swift
index 3b21738..0663683 100644
--- a/CourierE2EApp/Views/ConnectionView.swift
+++ b/CourierE2EApp/Views/ConnectionView.swift
@@ -85,7 +85,7 @@ struct ConnectionView: View {
                 HStack {
                     Text("QoS")
                     Picker(selection: $connectionVM.publishQoS) {
-                        ForEach(0..<3) { index in
+                        ForEach(0..<5) { index in
                             Text(String(index)).tag(index)
                         }
                     } label: {
diff --git a/CourierMQTT/MQTT/Connection/MQTT-Client-Framework/MQTTClientFrameworkSessionManager.swift b/CourierMQTT/MQTT/Connection/MQTT-Client-Framework/MQTTClientFrameworkSessionManager.swift
index dc2a5ad..a0f54b0 100644
--- a/CourierMQTT/MQTT/Connection/MQTT-Client-Framework/MQTTClientFrameworkSessionManager.swift
+++ b/CourierMQTT/MQTT/Connection/MQTT-Client-Framework/MQTTClientFrameworkSessionManager.swift
@@ -224,7 +224,27 @@ class MQTTClientFrameworkSessionManager: NSObject, IMQTTClientFrameworkSessionMa
 
     @discardableResult
     private func sendData(_ data: Data, topic: String, qos: MQTTQosLevel, retainFlag: Bool) -> UInt16 {
-        let msgId = session?.publishData(data, onTopic: topic, retain: retainFlag, qos: qos, publishHandler: nil) ?? 0
+        
+        var midProvider: (() -> UInt16)?
+        var publishHandler: MQTTPublishHandler?
+        if qos == MQTTQosLevel.atLeastOnceWithoutPersistenceAndNoRetry || qos == MQTTQosLevel.atLeastOnceWithoutPersistenceAndNoRetry {
+            publishHandler = { [weak self] error in
+                guard let self = self else { return }
+                printDebug("COURIER: Puback Handler for Special QoSes")
+                if error == nil {
+                    let mid = midProvider?() ?? 0
+                    self.delegate?.sessionManager(self, didDeliverMessageID: mid, topic: topic, data: data, qos: qos, retainFlag: retainFlag)
+                }
+            }
+        }
+        
+        let msgId = session?.publishData(data, onTopic: topic, retain: retainFlag, qos: qos, publishHandler: publishHandler) ?? 0
+        if publishHandler != nil {
+            midProvider = { [weak self] in
+                guard self != nil else { return 0 }
+                return msgId
+            }
+        }
         return msgId
     }
 
diff --git a/Internal/MQTT-Client-Framework/MQTTClientGJ/MQTTClientGJ/MQTTSession.m b/Internal/MQTT-Client-Framework/MQTTClientGJ/MQTTClientGJ/MQTTSession.m
index ead3443..729b8c3 100644
--- a/Internal/MQTT-Client-Framework/MQTTClientGJ/MQTTClientGJ/MQTTSession.m
+++ b/Internal/MQTT-Client-Framework/MQTTClientGJ/MQTTClientGJ/MQTTSession.m
@@ -401,6 +401,11 @@ - (UInt16)publishData:(NSData *)data
         if (qos == MQTTQosLevelAtLeastOnceWithoutPersistenceAndRetry || qos == MQTTQosLevelAtLeastOnceWithoutPersistenceAndNoRetry) {
             qos = MQTTQosLevelAtLeastOnce;
             msgId = [self nextMsgId];
+            if (publishHandler) {
+                (self.publishHandlers)[@(msgId)] = [publishHandler copy];
+            } else {
+                [self.publishHandlers removeObjectForKey:@(msgId)];
+            }
         }
         MQTTMessage *msg = [MQTTMessage publishMessageWithData:data
                                                        onTopic:topic
@@ -422,7 +427,7 @@ - (UInt16)publishData:(NSData *)data
                                         code:MQTTSessionErrorEncoderNotReady
                                     userInfo:@{NSLocalizedDescriptionKey : @"Encoder not ready"}];
         }
-        if (publishHandler) {
+        if (!qos && publishHandler) {
             [self onPublish:publishHandler error:error];
         }
     } else {
@@ -1161,6 +1166,12 @@ - (void)handlePuback:(MQTTMessage*)msg {
             [self.persistence sync];
             [self tell];
         }
+    } else {
+        MQTTPublishHandler publishHandler = (self.publishHandlers)[@(msg.mid)];
+        if (publishHandler) {
+            [self.publishHandlers removeObjectForKey:@(msg.mid)];
+            [self onPublish:publishHandler error:nil];
+        }
     }
 }