Remove handled fragmented messages from the map (#571)

* Remove handled fragmented messages from the map

We should remove handled fragmented messages from fragmentedMessages
map to free it up. This PR includes this fix.

Also, FragmentedClientMessageHandlerTest is enhanced to check the
the map's size before and after merging and handling the
end fragment.

* address review comments

Author: mdumandag

src/network/ClientConnection.ts

@@ -242,25 +242,32 @@ export class ClientMessageReader {
 /** @internal */
 export class FragmentedClientMessageHandler {
     private readonly fragmentedMessages = new Map<number, ClientMessage>();
+    private readonly logger: ILogger;
+
+    constructor(logger: ILogger) {
+        this.logger = logger;
+    }
 
     handleFragmentedMessage(clientMessage: ClientMessage, callback: Function): void {
         const fragmentationFrame = clientMessage.startFrame;
         const fragmentationId = clientMessage.getFragmentationId();
         clientMessage.dropFragmentationFrame();
         if (fragmentationFrame.hasBeginFragmentFlag()) {
             this.fragmentedMessages.set(fragmentationId, clientMessage);
-        } else if (fragmentationFrame.hasEndFragmentFlag()) {
-            const mergedMessage = this.mergeIntoExistingClientMessage(fragmentationId, clientMessage);
-            callback(mergedMessage);
         } else {
-            this.mergeIntoExistingClientMessage(fragmentationId, clientMessage);
-        }
-    }
+            const existingMessage = this.fragmentedMessages.get(fragmentationId);
+            if (existingMessage == null) {
+                this.logger.debug('FragmentedClientMessageHandler',
+                    'A fragmented message without the begin part is received. Fragmentation id: ' + fragmentationId);
+                return;
+            }
 
-    private mergeIntoExistingClientMessage(fragmentationId: number, clientMessage: ClientMessage): ClientMessage {
-        const existingMessage = this.fragmentedMessages.get(fragmentationId);
-        existingMessage.merge(clientMessage);
-        return existingMessage;
+            existingMessage.merge(clientMessage);
+            if (fragmentationFrame.hasEndFragmentFlag()) {
+                this.fragmentedMessages.delete(fragmentationId);
+                callback(existingMessage);
+            }
+        }
     }
 }
 
@@ -304,7 +311,7 @@ export class ClientConnection {
         this.reader = new ClientMessageReader();
         this.connectionId = connectionId;
         this.logger = this.client.getLoggingService().getLogger();
-        this.fragmentedMessageHandler = new FragmentedClientMessageHandler();
+        this.fragmentedMessageHandler = new FragmentedClientMessageHandler(this.logger);
     }
 
     /**

test/connection/FragmentedClientMessageHandlerTest.js

@@ -24,14 +24,16 @@ const ClientMessageReader = require('../../lib/network/ClientConnection').Client
 const cm = require('../../lib/ClientMessage');
 const FixSizedTypesCodec = require('../../lib/codec/builtin/FixSizedTypesCodec').FixSizedTypesCodec;
 
-describe('FragmentedClientMessageHandler', function () {
+describe('FragmentedClientMessageHandlerTest', function () {
 
     let reader;
     let handler;
 
     beforeEach(() => {
         reader = new ClientMessageReader();
-        handler = new FragmentedClientMessageHandler();
+        handler = new FragmentedClientMessageHandler({
+            debug: () => {}
+        });
     });
 
     it('handles fragmented message', function (done) {
@@ -62,15 +64,25 @@ describe('FragmentedClientMessageHandler', function () {
 
         // Should merge with the above message
         handler.handleFragmentedMessage(fragment1, () => done(new Error("It should just merge.")));
+        expect(handler.fragmentedMessages.size).to.equal(1);
         compareMessages(fragments[1].startFrame.next, endFrame.next);
         endFrame = fragmentedMessage.endFrame;
 
         // Should merge with the above message and run the callback
+        let isCalled = false;
         handler.handleFragmentedMessage(fragment2, () => {
             compareMessages(fragments[2].startFrame.next, endFrame.next);
-            done();
+            isCalled = true;
         });
+        expect(isCalled).to.be.true;
+        expect(handler.fragmentedMessages.size).to.equal(0);
+
+        // If a message with a missing begin part is received, we should do nothing
+        handler.handleFragmentedMessage(fragment1, () => done(new Error("It should ignore invalid messages.")))
+        handler.handleFragmentedMessage(fragment2, () => done(new Error("It should ignore invalid messages.")))
+        expect(handler.fragmentedMessages.size).to.equal(0);
 
+        done();
     });
 
     const compareMessages = (frame1, frame2) => {