From 633688f184aa1f11416acbeb6bcca74f795cbd9d Mon Sep 17 00:00:00 2001
From: "selva.andre" <selva.andre@gmail.com>
Date: Tue, 11 Sep 2012 23:20:32 +0200
Subject: [PATCH] Implemented the QoS 2 sending part for the server plus  a
 unit test

---
 .../messaging/spi/impl/SimpleMessaging.java   | 29 ++++++++++++++++++
 .../dna/mqtt/moquette/server/MQTTHandler.java |  3 ++
 .../server/ServerIntegrationFuseTest.java     | 30 +++++++++++++++++++
 3 files changed, 62 insertions(+)

diff --git a/parser/src/main/java/org/dna/mqtt/moquette/messaging/spi/impl/SimpleMessaging.java b/parser/src/main/java/org/dna/mqtt/moquette/messaging/spi/impl/SimpleMessaging.java
index 9336c689..f8ea68c0 100644
--- a/parser/src/main/java/org/dna/mqtt/moquette/messaging/spi/impl/SimpleMessaging.java
+++ b/parser/src/main/java/org/dna/mqtt/moquette/messaging/spi/impl/SimpleMessaging.java
@@ -183,6 +183,14 @@ public class SimpleMessaging implements IMessaging, EventHandler<ValueEvent> {
                 String clientID = (String) session.getAttribute(Constants.ATTR_CLIENTID);
                 int messageID = ((PubRelMessage) message).getMessageID();
                 processPubRel(clientID, messageID);
+            } else if (message instanceof PubRecMessage) {
+                String clientID = (String) session.getAttribute(Constants.ATTR_CLIENTID);
+                int messageID = ((PubRecMessage) message).getMessageID();
+                processPubRec(clientID, messageID);
+            } else if (message instanceof PubCompMessage) {
+                String clientID = (String) session.getAttribute(Constants.ATTR_CLIENTID);
+                int messageID = ((PubCompMessage) message).getMessageID();
+                processPubComp(clientID, messageID);
             } else {
                 throw new RuntimeException("Illegal message received " + message);
             }
@@ -351,12 +359,33 @@ public class SimpleMessaging implements IMessaging, EventHandler<ValueEvent> {
                     PublishEvent newPublishEvt = new PublishEvent(topic, qos, message, retain, sub.getClientId(), messageID, null);
                     m_storageService.storePublishForFuture(newPublishEvt);
                 } else  {
+                    //if QoS 2 then store it in temp memory
+                    if (qos ==QOSType.EXACTLY_ONCE) {
+                        String publishKey = String.format("%s%d", sub.getClientId(), messageID);
+                        PublishEvent newPublishEvt = new PublishEvent(topic, qos, message, retain, sub.getClientId(), messageID, null);
+                        m_storageService.addInFlight(newPublishEvt, publishKey);
+                    }
+                    //publish
                     notify(new NotifyEvent(sub.getClientId(), topic, qos, message, false));
                 }
             }
         }
     }
 
+    private void processPubRec(String clientID, int messageID) {
+        //once received a PUBREC reply with a PUBREL(messageID)
+        LOG.debug(String.format("processPubRec invoked for clientID %s ad messageID %d", clientID, messageID));
+        PubRelMessage pubRelMessage = new PubRelMessage();
+        pubRelMessage.setMessageID(messageID);
+
+        m_clientIDs.get(clientID).getSession().write(pubRelMessage);
+    }
+
+    private void processPubComp(String clientID, int messageID) {
+        //once received the PUBCOMP then remove the message from the temp memory
+        String publishKey = String.format("%s%d", clientID, messageID);
+        m_storageService.cleanInFlight(publishKey);
+    }
 
     private void subscribeSingleTopic(Subscription newSubscription, final String topic) {
         subscriptions.add(newSubscription);
diff --git a/parser/src/main/java/org/dna/mqtt/moquette/server/MQTTHandler.java b/parser/src/main/java/org/dna/mqtt/moquette/server/MQTTHandler.java
index 76d3892a..7bdc07fe 100644
--- a/parser/src/main/java/org/dna/mqtt/moquette/server/MQTTHandler.java
+++ b/parser/src/main/java/org/dna/mqtt/moquette/server/MQTTHandler.java
@@ -31,6 +31,9 @@ public class MQTTHandler extends IoHandlerAdapter {
                 case SUBSCRIBE:
                 case UNSUBSCRIBE:
                 case PUBLISH:
+                case PUBREC:
+                case PUBCOMP:
+                case PUBREL:
                 case PINGREQ:
                 case DISCONNECT:
                     m_messaging.handleProtocolMessage(session, msg);
diff --git a/parser/src/test/java/org/dna/mqtt/moquette/server/ServerIntegrationFuseTest.java b/parser/src/test/java/org/dna/mqtt/moquette/server/ServerIntegrationFuseTest.java
index 1789eb4c..71286d8f 100644
--- a/parser/src/test/java/org/dna/mqtt/moquette/server/ServerIntegrationFuseTest.java
+++ b/parser/src/test/java/org/dna/mqtt/moquette/server/ServerIntegrationFuseTest.java
@@ -145,4 +145,34 @@ public class ServerIntegrationFuseTest {
 
         assertEquals("Hello MQTT", new String(message.getPayload()));
     }
+
+    @Test
+    public void testPublishReceiveWithQoS2() throws Exception {
+        m_mqtt.setCleanSession(false);
+        m_connection = m_mqtt.futureConnection();
+        m_connection.connect().await();
+
+        //subscribe to a QoS 2 topic
+        Topic[] topics = {new Topic("/topic", QoS.EXACTLY_ONCE)};
+        Future<byte[]> subscribeFuture = m_connection.subscribe(topics);
+        byte[] qoses = subscribeFuture.await();
+
+//        m_connection.disconnect().await();
+
+        //publish a QoS 1 message another client publish a message on the topic
+        publishFromAnotherClient("/topic", "Hello MQTT".getBytes(), QoS.EXACTLY_ONCE);
+
+
+        //Reconnect
+        /*m_connection = m_mqtt.futureConnection();
+        m_connection.connect().await();*/
+        Future<Message> receive = m_connection.receive();
+
+        //verify the reception
+        Message message = receive.await();
+        message.ack();
+
+        assertEquals("Hello MQTT", new String(message.getPayload()));
+    }
+
 }