From 82f6736bce1413cabdb06e248e86f164a9fbf7b4 Mon Sep 17 00:00:00 2001
From: warunalakshitha <warudanu@gmail.com>
Date: Sun, 1 Oct 2017 10:57:26 +0530
Subject: [PATCH] fix concurrent accesss issue with remote session

---
 .../pom.xml                                   |   2 +-
 .../ClientSessionSubscriptionEndpoint.java    |  33 +--
 .../DeviceSessionSubscriptionEndpoint.java    |  30 +--
 .../endpoint/SubscriptionEndpoint.java        |  44 ++--
 .../session/endpoint/constants/Constants.java |   3 +
 .../utils/HttpSessionConfigurator.java        |   2 +-
 .../pom.xml                                   |   4 +-
 .../RemoteSessionManagementService.java       |  25 +-
 .../RemoteSessionManagementServiceImpl.java   | 226 ++++++++++--------
 .../authentication/AuthenticationInfo.java    |   2 +-
 .../authentication/OAuthAuthenticator.java    |   2 +-
 .../oauth/OAuthTokenValidator.java            |   2 +-
 .../oauth/OAuthTokenValidatorStubFactory.java |   3 +-
 .../OAuthTokenValidationException.java        |   2 +-
 .../constants/RemoteSessionConstants.java     |  12 +-
 .../remote.session/dto/RemoteSession.java     |  66 ++---
 .../RemoteSessionManagementException.java     |  21 +-
 .../RemoteSessionManagementDataHolder.java    |  31 +--
 ...moteSessionManagementServiceComponent.java |  23 +-
 .../RemoteSessionManagerStartupListener.java  |   8 +-
 .../remote.session/util/PropertyUtils.java    |  12 +-
 .../remote-session-extension/pom.xml          |   2 +-
 22 files changed, 268 insertions(+), 287 deletions(-)

diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/pom.xml b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/pom.xml
index a2b67a2933..88edfc0996 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/pom.xml
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/pom.xml
@@ -23,7 +23,7 @@
     <parent>
         <groupId>org.wso2.carbon.devicemgt-plugins</groupId>
         <artifactId>remote-session-extension</artifactId>
-        <version>4.0.54-SNAPSHOT</version>
+        <version>4.0.66-SNAPSHOT</version>
         <relativePath>../pom.xml</relativePath>
     </parent>
 
diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/src/main/java/org/wso2/carbon/device/mgt/extensions/remote/session/endpoint/ClientSessionSubscriptionEndpoint.java b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/src/main/java/org/wso2/carbon/device/mgt/extensions/remote/session/endpoint/ClientSessionSubscriptionEndpoint.java
index 66ab3e89f1..cadb860886 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/src/main/java/org/wso2/carbon/device/mgt/extensions/remote/session/endpoint/ClientSessionSubscriptionEndpoint.java
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/src/main/java/org/wso2/carbon/device/mgt/extensions/remote/session/endpoint/ClientSessionSubscriptionEndpoint.java
@@ -20,11 +20,9 @@ package org.wso2.carbon.device.mgt.extensions.remote.session.endpoint;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
-import org.wso2.carbon.device.mgt.extensions.remote.session.constants.RemoteSessionConstants;
-import org.wso2.carbon.device.mgt.extensions.remote.session.exception.RemoteSessionInvalidException;
-import org.wso2.carbon.device.mgt.extensions.remote.session.exception.RemoteSessionManagementException;
 import org.wso2.carbon.device.mgt.extensions.remote.session.endpoint.utils.HttpSessionConfigurator;
 import org.wso2.carbon.device.mgt.extensions.remote.session.endpoint.utils.ServiceHolder;
+import org.wso2.carbon.device.mgt.extensions.remote.session.exception.RemoteSessionManagementException;
 
 import javax.websocket.CloseReason;
 import javax.websocket.OnClose;
@@ -43,6 +41,7 @@ import java.io.IOException;
 public class ClientSessionSubscriptionEndpoint extends SubscriptionEndpoint {
 
     private static final Log log = LogFactory.getLog(ClientSessionSubscriptionEndpoint.class);
+
     /**
      * Web socket onOpen - When client sends a message
      *
@@ -51,15 +50,10 @@ public class ClientSessionSubscriptionEndpoint extends SubscriptionEndpoint {
     @OnOpen
     public void onOpen(Session session, @PathParam("deviceType") String deviceType, @PathParam("deviceId") String
             deviceId) {
-        System.out.print("**************Open***************");
-        if (log.isDebugEnabled()) {
-            log.debug("WebSocket opened, for RemoteSession id: " + session.getId());
-        }
         try {
-            ServiceHolder.getInstance().getRemoteSessionManagementService().initializeSession(session, deviceType, deviceId);
-            System.out.print("**************Opened***************");
-        } catch (RemoteSessionInvalidException e) {
-            System.out.print(e.getMessage());
+            ServiceHolder.getInstance().getRemoteSessionManagementService().initializeSession(session, deviceType,
+                    deviceId);
+        } catch (RemoteSessionManagementException e) {
             if (log.isDebugEnabled()) {
                 log.error("Error occurred while initializing session ", e);
             }
@@ -68,20 +62,7 @@ public class ClientSessionSubscriptionEndpoint extends SubscriptionEndpoint {
             } catch (IOException ex) {
                 log.error("Failed to disconnect the client.", ex);
             }
-        } catch (RemoteSessionManagementException e) {
-            System.out.print(e.getMessage());
-            if (log.isDebugEnabled()) {
-                log.error("Error occurred while initializing session ", e);
-            }
-            try {
-                session.close(new CloseReason(CloseReason.CloseCodes.CANNOT_ACCEPT, "Error occurred while adding operation"));
-            } catch (IOException ex) {
-                if (log.isDebugEnabled()) {
-                    log.error("Failed to disconnect the client.", ex);
-                }
-            }
         }
-
     }
 
     /**
@@ -92,7 +73,7 @@ public class ClientSessionSubscriptionEndpoint extends SubscriptionEndpoint {
      */
     @OnMessage
     public void onMessage(Session session, String message, @PathParam("deviceType") String deviceType, @PathParam
-            ("deviceId") String deviceId) throws RemoteSessionManagementException {
+            ("deviceId") String deviceId) {
         super.onMessage(session, message, deviceType, deviceId);
     }
 
@@ -104,7 +85,7 @@ public class ClientSessionSubscriptionEndpoint extends SubscriptionEndpoint {
      */
     @OnMessage
     public void onMessage(Session session, byte[] message, @PathParam("deviceType") String deviceType, @PathParam
-            ("deviceId") String deviceId) throws RemoteSessionManagementException {
+            ("deviceId") String deviceId) {
         super.onMessage(session, message, deviceType, deviceId);
     }
 
diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/src/main/java/org/wso2/carbon/device/mgt/extensions/remote/session/endpoint/DeviceSessionSubscriptionEndpoint.java b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/src/main/java/org/wso2/carbon/device/mgt/extensions/remote/session/endpoint/DeviceSessionSubscriptionEndpoint.java
index 9c771f74d1..9933f1f616 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/src/main/java/org/wso2/carbon/device/mgt/extensions/remote/session/endpoint/DeviceSessionSubscriptionEndpoint.java
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/src/main/java/org/wso2/carbon/device/mgt/extensions/remote/session/endpoint/DeviceSessionSubscriptionEndpoint.java
@@ -20,11 +20,9 @@ package org.wso2.carbon.device.mgt.extensions.remote.session.endpoint;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
-import org.wso2.carbon.device.mgt.extensions.remote.session.constants.RemoteSessionConstants;
-import org.wso2.carbon.device.mgt.extensions.remote.session.exception.RemoteSessionInvalidException;
-import org.wso2.carbon.device.mgt.extensions.remote.session.exception.RemoteSessionManagementException;
 import org.wso2.carbon.device.mgt.extensions.remote.session.endpoint.utils.HttpSessionConfigurator;
 import org.wso2.carbon.device.mgt.extensions.remote.session.endpoint.utils.ServiceHolder;
+import org.wso2.carbon.device.mgt.extensions.remote.session.exception.RemoteSessionManagementException;
 
 import javax.websocket.CloseReason;
 import javax.websocket.OnClose;
@@ -43,6 +41,7 @@ import java.io.IOException;
 public class DeviceSessionSubscriptionEndpoint extends SubscriptionEndpoint {
 
     private static final Log log = LogFactory.getLog(DeviceSessionSubscriptionEndpoint.class);
+
     /**
      * Web socket onOpen - When client sends a message
      *
@@ -51,16 +50,10 @@ public class DeviceSessionSubscriptionEndpoint extends SubscriptionEndpoint {
     @OnOpen
     public void onOpen(Session session, @PathParam("deviceType") String deviceType, @PathParam("deviceId") String
             deviceId, @PathParam("operationId") String operationId) {
-        System.out.print("**************Open***************"+operationId);
-        if (log.isDebugEnabled()) {
-            log.debug("WebSocket opened, for RemoteSession id: " + session.getId());
-        }
         try {
             ServiceHolder.getInstance().getRemoteSessionManagementService().initializeSession(session, deviceType,
                     deviceId, operationId);
-            System.out.print("**************Opened***************");
-        } catch (RemoteSessionInvalidException e) {
-            System.out.print(e.getMessage());
+        } catch (RemoteSessionManagementException e) {
             if (log.isDebugEnabled()) {
                 log.error("Error occurred while initializing session ", e);
             }
@@ -69,20 +62,7 @@ public class DeviceSessionSubscriptionEndpoint extends SubscriptionEndpoint {
             } catch (IOException ex) {
                 log.error("Failed to disconnect the client.", ex);
             }
-        } catch (RemoteSessionManagementException e) {
-            System.out.print(e.getMessage());
-            if (log.isDebugEnabled()) {
-                log.error("Error occurred while initializing session ", e);
-            }
-            try {
-                session.close(new CloseReason(CloseReason.CloseCodes.CANNOT_ACCEPT, "Error occurred while adding operation"));
-            } catch (IOException ex) {
-                if (log.isDebugEnabled()) {
-                    log.error("Failed to disconnect the client.", ex);
-                }
-            }
         }
-
     }
 
     /**
@@ -93,7 +73,7 @@ public class DeviceSessionSubscriptionEndpoint extends SubscriptionEndpoint {
      */
     @OnMessage
     public void onMessage(Session session, String message, @PathParam("deviceType") String deviceType, @PathParam
-            ("deviceId") String deviceId) throws RemoteSessionManagementException {
+            ("deviceId") String deviceId) {
         super.onMessage(session, message, deviceType, deviceId);
     }
 
@@ -105,7 +85,7 @@ public class DeviceSessionSubscriptionEndpoint extends SubscriptionEndpoint {
      */
     @OnMessage
     public void onMessage(Session session, byte[] message, @PathParam("deviceType") String deviceType, @PathParam
-            ("deviceId") String deviceId) throws RemoteSessionManagementException {
+            ("deviceId") String deviceId) {
         super.onMessage(session, message, deviceType, deviceId);
     }
 
diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/src/main/java/org/wso2/carbon/device/mgt/extensions/remote/session/endpoint/SubscriptionEndpoint.java b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/src/main/java/org/wso2/carbon/device/mgt/extensions/remote/session/endpoint/SubscriptionEndpoint.java
index 53cd4def73..2431a674ef 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/src/main/java/org/wso2/carbon/device/mgt/extensions/remote/session/endpoint/SubscriptionEndpoint.java
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/src/main/java/org/wso2/carbon/device/mgt/extensions/remote/session/endpoint/SubscriptionEndpoint.java
@@ -20,9 +20,8 @@ package org.wso2.carbon.device.mgt.extensions.remote.session.endpoint;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
-import org.wso2.carbon.device.mgt.extensions.remote.session.exception.RemoteSessionInvalidException;
-import org.wso2.carbon.device.mgt.extensions.remote.session.exception.RemoteSessionManagementException;
 import org.wso2.carbon.device.mgt.extensions.remote.session.endpoint.utils.ServiceHolder;
+import org.wso2.carbon.device.mgt.extensions.remote.session.exception.RemoteSessionManagementException;
 
 import javax.websocket.CloseReason;
 import javax.websocket.Session;
@@ -43,15 +42,14 @@ public class SubscriptionEndpoint {
      * @param message - Status code for web-socket close.
      */
     public void onMessage(Session session, String message, @PathParam("deviceType") String deviceType, @PathParam
-            ("deviceId") String deviceId) throws RemoteSessionManagementException {
-        System.out.print("______________" + session.getId());
+            ("deviceId") String deviceId) {
         if (log.isDebugEnabled()) {
             log.debug("Received message from client for RemoteSession id: " + session.getId() + " device type: " +
                     deviceType + " device id: " + deviceId);
         }
         try {
             ServiceHolder.getInstance().getRemoteSessionManagementService().sendMessageToPeer(session, message);
-        } catch (RemoteSessionInvalidException e) {
+        } catch (RemoteSessionManagementException e) {
             if (log.isDebugEnabled()) {
                 log.error("Error occurred while send message to peer session ", e);
             }
@@ -72,15 +70,14 @@ public class SubscriptionEndpoint {
      * @param message - Message which needs to send to peer
      */
     public void onMessage(Session session, byte[] message, @PathParam("deviceType") String deviceType, @PathParam
-            ("deviceId") String deviceId) throws RemoteSessionManagementException {
-        System.out.print("______________" + session.getId());
+            ("deviceId") String deviceId) {
         if (log.isDebugEnabled()) {
             log.debug("Received message from client for RemoteSession id: " + session.getId() + " device type: " +
                     deviceType + " device id: " + deviceId);
         }
         try {
             ServiceHolder.getInstance().getRemoteSessionManagementService().sendMessageToPeer(session, message);
-        } catch (RemoteSessionInvalidException e) {
+        } catch (RemoteSessionManagementException e) {
             if (log.isDebugEnabled()) {
                 log.error("Error occurred while send message to peer session ", e);
             }
@@ -102,17 +99,14 @@ public class SubscriptionEndpoint {
      */
     public void onClose(Session session, CloseReason reason, @PathParam("deviceType") String deviceType, @PathParam
             ("deviceId") String deviceId) {
+
+        ServiceHolder.getInstance().getRemoteSessionManagementService().endSession(session, "Remote session closed");
         if (log.isDebugEnabled()) {
-            log.debug("Closing a WebSocket due to " + reason.getReasonPhrase() + ", for session ID:" + session.getId
-                    () + ", for request URI - " + session.getRequestURI() + " device type: " + deviceType + " device id: " + deviceId);
-        }
-        try {
-            ServiceHolder.getInstance().getRemoteSessionManagementService().endSession(session);
-        } catch (IOException ex) {
-            if (log.isDebugEnabled()) {
-                log.error("Failed to disconnect the client.", ex);
-            }
+            log.debug("websocket closed due to " + reason.getReasonPhrase() + ", for session ID:" + session.getId
+                    () + ", for request URI - " + session.getRequestURI() + " device type: " + deviceType + " device " +
+                    "id: " + deviceId);
         }
+
     }
 
     /**
@@ -123,12 +117,20 @@ public class SubscriptionEndpoint {
      */
     public void onError(Session session, Throwable throwable, @PathParam("deviceType") String deviceType, @PathParam
             ("deviceId") String deviceId) {
-        log.error(
-                "Error occurred in session ID: " + session.getId() + " device type: " + deviceType + " device id: " +
-                        deviceId + ", for request URI - " + session.getRequestURI() +
+
+        if (throwable instanceof IOException) {
+            if (log.isDebugEnabled()) {
+                log.error("Error occurred in session ID: " + session.getId() + " device type: " + deviceType +
+                        "device id: " + deviceId + ", for request URI - " + session.getRequestURI() +
                         ", " + throwable.getMessage(), throwable);
+            }
+        } else {
+            log.error("Error occurred in session ID: " + session.getId() + " device type: " + deviceType + " device " +
+                    "id: " + deviceId + ", for request URI - " + session.getRequestURI() + ", " + throwable.getMessage
+                    (), throwable);
+        }
         try {
-            ServiceHolder.getInstance().getRemoteSessionManagementService().endSession(session);
+            ServiceHolder.getInstance().getRemoteSessionManagementService().endSession(session, "Remote session closed");
             if (session.isOpen()) {
                 session.close(new CloseReason(CloseReason.CloseCodes.PROTOCOL_ERROR, "Unexpected Error Occurred"));
             }
diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/src/main/java/org/wso2/carbon/device/mgt/extensions/remote/session/endpoint/constants/Constants.java b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/src/main/java/org/wso2/carbon/device/mgt/extensions/remote/session/endpoint/constants/Constants.java
index 9bf277389b..fc3dcdfbcc 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/src/main/java/org/wso2/carbon/device/mgt/extensions/remote/session/endpoint/constants/Constants.java
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/src/main/java/org/wso2/carbon/device/mgt/extensions/remote/session/endpoint/constants/Constants.java
@@ -20,6 +20,9 @@
 
 package org.wso2.carbon.device.mgt.extensions.remote.session.endpoint.constants;
 
+/**
+ * This holds the constants related to remote session web socket endpoint
+ */
 public class Constants {
     public static final String HTTP_HEADERS = "HttpHeaders";
 }
diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/src/main/java/org/wso2/carbon/device/mgt/extensions/remote/session/endpoint/utils/HttpSessionConfigurator.java b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/src/main/java/org/wso2/carbon/device/mgt/extensions/remote/session/endpoint/utils/HttpSessionConfigurator.java
index a0088689b4..94b6bf3fd3 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/src/main/java/org/wso2/carbon/device/mgt/extensions/remote/session/endpoint/utils/HttpSessionConfigurator.java
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session.endpoint/src/main/java/org/wso2/carbon/device/mgt/extensions/remote/session/endpoint/utils/HttpSessionConfigurator.java
@@ -1,6 +1,6 @@
 /*
  *
- *  Copyright (c) 2016, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
+ *  Copyright (c) 2017, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
  *
  *  WSO2 Inc. licenses this file to you under the Apache License,
  *  Version 2.0 (the "License"); you may not use this file except
diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/pom.xml b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/pom.xml
index a8033555f7..659aac91f6 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/pom.xml
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/pom.xml
@@ -1,6 +1,6 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <!--
-  ~ Copyright (c) 2016, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
+  ~ Copyright (c) 2017, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
   ~
   ~ WSO2 Inc. licenses this file to you under the Apache License,
   ~ Version 2.0 (the "License"); you may not use this file except
@@ -22,7 +22,7 @@
     <parent>
         <groupId>org.wso2.carbon.devicemgt-plugins</groupId>
         <artifactId>remote-session-extension</artifactId>
-        <version>4.0.54-SNAPSHOT</version>
+        <version>4.0.66-SNAPSHOT</version>
         <relativePath>../pom.xml</relativePath>
     </parent>
 
diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/RemoteSessionManagementService.java b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/RemoteSessionManagementService.java
index 0cff39b89d..53db138812 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/RemoteSessionManagementService.java
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/RemoteSessionManagementService.java
@@ -17,7 +17,6 @@
  */
 package org.wso2.carbon.device.mgt.extensions.remote.session;
 
-import org.wso2.carbon.device.mgt.extensions.remote.session.exception.RemoteSessionInvalidException;
 import org.wso2.carbon.device.mgt.extensions.remote.session.exception.RemoteSessionManagementException;
 
 import javax.websocket.Session;
@@ -34,11 +33,10 @@ public interface RemoteSessionManagementService {
      * @param session    Web socket RemoteSession
      * @param deviceType Device Type
      * @param deviceId   Device Id
-     * @throws RemoteSessionInvalidException    throws when session cannot be made due to invalid data
-     * @throws RemoteSessionManagementException throws when session has error with accessing device resources
+     * @throws RemoteSessionManagementException throws when session has errors with accessing device resources
      */
     public void initializeSession(Session session, String deviceType, String deviceId) throws
-            RemoteSessionInvalidException, RemoteSessionManagementException;
+            RemoteSessionManagementException;
 
     /**
      * Initialize session based on web socket request . This method use by the device to connect
@@ -47,39 +45,34 @@ public interface RemoteSessionManagementService {
      * @param deviceType  Device Type
      * @param deviceId    Device Id
      * @param operationId Operation Id that device needs to connec
-     * @throws RemoteSessionInvalidException    throws when session cannot be made due to invalid data
-     * @throws RemoteSessionManagementException throws when session has error with accessing device resources
+     * @throws RemoteSessionManagementException throws when session has errors with accessing device resources
      */
     public void initializeSession(Session session, String deviceType, String deviceId, String operationId) throws
-            RemoteSessionInvalidException, RemoteSessionManagementException;
+            RemoteSessionManagementException;
 
     /**
-     * Send message to connected remote device or client
+     * Send string message to connected remote device or client
      *
      * @param session Web socket RemoteSession
      * @param message Message needs to send to peer connection client
-     * @throws RemoteSessionInvalidException
      * @throws RemoteSessionManagementException
      */
-    public void sendMessageToPeer(Session session, String message) throws
-            RemoteSessionInvalidException, RemoteSessionManagementException;
+    public void sendMessageToPeer(Session session, String message) throws  RemoteSessionManagementException;
 
     /**
-     * Send message to connected remote device or client
+     * Send byte message to connected remote device or client
      *
      * @param session Web socket RemoteSession
      * @param message Message needs to send to peer connection
-     * @throws RemoteSessionInvalidException
      * @throws RemoteSessionManagementException
      */
-    public void sendMessageToPeer(Session session, byte[] message) throws
-            RemoteSessionInvalidException, RemoteSessionManagementException;
+    public void sendMessageToPeer(Session session, byte[] message) throws RemoteSessionManagementException;
 
     /**
      * Close the session
      *
      * @param session Web socket RemoteSession
      */
-    public void endSession(Session session) throws IOException;
+    public void endSession(Session session,String closeReason);
 
 }
diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/RemoteSessionManagementServiceImpl.java b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/RemoteSessionManagementServiceImpl.java
index 437452ae7d..8c1a8a7ab3 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/RemoteSessionManagementServiceImpl.java
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/RemoteSessionManagementServiceImpl.java
@@ -33,10 +33,7 @@ import org.wso2.carbon.device.mgt.core.operation.mgt.ConfigOperation;
 import org.wso2.carbon.device.mgt.extensions.remote.session.authentication.AuthenticationInfo;
 import org.wso2.carbon.device.mgt.extensions.remote.session.authentication.OAuthAuthenticator;
 import org.wso2.carbon.device.mgt.extensions.remote.session.constants.RemoteSessionConstants;
-import org.wso2.carbon.device.mgt.extensions.remote.session.dto.ClientSession;
-import org.wso2.carbon.device.mgt.extensions.remote.session.dto.DeviceSession;
-import org.wso2.carbon.device.mgt.extensions.remote.session.dto.common.RemoteSession;
-import org.wso2.carbon.device.mgt.extensions.remote.session.exception.RemoteSessionInvalidException;
+import org.wso2.carbon.device.mgt.extensions.remote.session.dto.RemoteSession;
 import org.wso2.carbon.device.mgt.extensions.remote.session.exception.RemoteSessionManagementException;
 import org.wso2.carbon.device.mgt.extensions.remote.session.internal.RemoteSessionManagementDataHolder;
 
@@ -53,21 +50,19 @@ import java.util.Map;
 public class RemoteSessionManagementServiceImpl implements RemoteSessionManagementService {
 
     private static final Log log = LogFactory.getLog(RemoteSessionManagementServiceImpl.class);
-    private static final String DATE_FORMAT_NOW = "yyyy-MM-dd HH:mm:ss";
-    private static final int MAX_BUFFER_SIZE = 640 * 1024;
-
 
     @Override
     public void initializeSession(Session session, String deviceType, String deviceId, String operationId) throws
-            RemoteSessionInvalidException, RemoteSessionManagementException {
-
+            RemoteSessionManagementException {
 
+        // Check whether required configurations are enabled
         if (!RemoteSessionManagementDataHolder.getInstance().isEnabled()) {
             throw new RemoteSessionManagementException("Remote session feature is disabled.");
         } else if (RemoteSessionManagementDataHolder.getInstance().getServerUrl() == null) {
-            throw new RemoteSessionManagementException("Server url haven't been configured.");
+            throw new RemoteSessionManagementException("Server url has not been configured.");
         }
 
+        // Read Query Parameters for obtain the token
         Map<String, List<String>> sessionQueryParam = new HashedMap();
         List<String> sessionQueryParamList = new LinkedList<>();
         sessionQueryParamList.add(session.getQueryString());
@@ -78,64 +73,61 @@ public class RemoteSessionManagementServiceImpl implements RemoteSessionManageme
         AuthenticationInfo authenticationInfo = oAuthAuthenticator.isAuthenticated(sessionQueryParam);
 
         if (authenticationInfo != null && authenticationInfo.isAuthenticated()) {
-
             try {
                 PrivilegedCarbonContext.startTenantFlow();
-                PrivilegedCarbonContext.getThreadLocalCarbonContext().setTenantDomain(authenticationInfo.getTenantDomain()
+                PrivilegedCarbonContext.getThreadLocalCarbonContext().setTenantDomain(authenticationInfo
+                                .getTenantDomain()
                         , true);
                 PrivilegedCarbonContext.getThreadLocalCarbonContext().setUsername(authenticationInfo.getUsername());
                 if (deviceId != null && !deviceId.isEmpty() && deviceType != null && !deviceType.isEmpty()) {
                     DeviceIdentifier deviceIdentifier = new DeviceIdentifier();
                     deviceIdentifier.setId(deviceId);
                     deviceIdentifier.setType(deviceType);
-                    // Check authorization for user
+
+                    // Check authorization of user for given device
                     boolean userAuthorized = RemoteSessionManagementDataHolder.getInstance()
                             .getDeviceAccessAuthorizationService()
                             .isUserAuthorized(deviceIdentifier, authenticationInfo.getUsername());
                     if (userAuthorized) {
-                        log.info("Operation ID: " + operationId);
-
                         // set common settings for session
-                        session.setMaxBinaryMessageBufferSize(MAX_BUFFER_SIZE);
-                        session.setMaxTextMessageBufferSize(MAX_BUFFER_SIZE);
+                        session.setMaxBinaryMessageBufferSize(RemoteSessionConstants.MAX_BUFFER_SIZE);
+                        session.setMaxTextMessageBufferSize(RemoteSessionConstants.MAX_BUFFER_SIZE);
                         session.setMaxIdleTimeout(RemoteSessionManagementDataHolder.getInstance().getMaxIdleTimeout());
 
-                        // if session initiated using operatiod id means request came from device
+                        // if session initiated using operation id means request came from device
                         if (operationId != null) {
-                            Session pendingSession = RemoteSessionManagementDataHolder.getInstance()
-                                    .getDeviceRequestMap().get((authenticationInfo.getTenantDomain() + "/" + deviceType
-                                            + "/" + deviceId));
-
-                            if (pendingSession != null) {
+                            RemoteSession activeSession = RemoteSessionManagementDataHolder.getInstance()
+                                    .getActiveDeviceClientSessionMap().get((authenticationInfo.getTenantDomain() + "/" +
+                                            deviceType + "/" + deviceId));
+                            if (activeSession != null) {
                                 RemoteSession clientRemote = RemoteSessionManagementDataHolder.getInstance()
-                                        .getSessionMap().get(pendingSession.getId());
+                                        .getSessionMap().get(activeSession.getMySession().getId());
                                 if (clientRemote != null) {
-
                                     if (clientRemote.getOperationId().equals(operationId)) {
-                                        RemoteSession deviceRemote = new DeviceSession(session, authenticationInfo
-                                                .getTenantDomain(), deviceType, deviceId, operationId);
+                                        RemoteSession deviceRemote = new RemoteSession(session, authenticationInfo
+                                                .getTenantDomain(), deviceType, deviceId);
+                                        deviceRemote.setOperationId(operationId);
                                         deviceRemote.setPeerSession(clientRemote);
                                         clientRemote.setPeerSession(deviceRemote);
                                         RemoteSessionManagementDataHolder.getInstance().getSessionMap().put(session
                                                 .getId(), deviceRemote);
-                                        RemoteSessionManagementDataHolder.getInstance().getDeviceRequestMap().remove(
-                                                (authenticationInfo.getTenantDomain() + "/" + deviceType + "/" + deviceId));
                                         // Send Remote connect response
                                         JSONObject message = new JSONObject();
-                                        message.put("code", RemoteSessionConstants.REMOTE_CONNECT);
-                                        message.put("operation_response", "connected");
+                                        message.put(RemoteSessionConstants.REMOTE_CONNECT_CODE, RemoteSessionConstants
+                                                .REMOTE_CONNECT);
                                         deviceRemote.sendMessageToPeer(message.toString());
-
+                                        log.info("Device session opened for session id: " + session.getId() +
+                                                " device Type : " + deviceType + " , " + "deviceId : " + deviceId);
                                     } else {
-                                        throw new RemoteSessionManagementException("Device and Operation information does" +
-                                                " not matched with client information for operation id: " + operationId +
-                                                " device Type : " + deviceType + " , " + "deviceId : " +
+                                        throw new RemoteSessionManagementException("Device and Operation information " +
+                                                "does not matched with client information for operation id: " +
+                                                operationId + " device Type : " + deviceType + " , " + "deviceId : " +
                                                 deviceId);
                                     }
                                 } else {
-                                    throw new RemoteSessionManagementException("Device session is inactive for operation " +
-                                            "id: " + operationId + " device Type : " + deviceType + " , " + "deviceId : " +
-                                            deviceId);
+                                    throw new RemoteSessionManagementException("Device session is inactive for " +
+                                            "operation id: " + operationId + " device Type : " + deviceType + " , " +
+                                            "deviceId : " + deviceId);
                                 }
 
 
@@ -144,19 +136,34 @@ public class RemoteSessionManagementServiceImpl implements RemoteSessionManageme
                                         "id: " + operationId + " device Type : " + deviceType + " , " + "deviceId : " +
                                         deviceId);
                             }
-
                         } else {
+                            RemoteSession clientRemote = new RemoteSession(session, authenticationInfo
+                                    .getTenantDomain(), deviceType, deviceId);
                             // Create new remote control operation to start the session
-                            Session pendingSession = RemoteSessionManagementDataHolder.getInstance().getDeviceRequestMap().get(
-                                    (authenticationInfo.getTenantDomain() + "/" + deviceType + "/" + deviceId));
-                            if (pendingSession != null && pendingSession.isOpen()) {
-                                throw new RemoteSessionManagementException("Another client session waiting on device to connect.");
+                            RemoteSession activeSession = RemoteSessionManagementDataHolder.getInstance()
+                                    .getActiveDeviceClientSessionMap().putIfAbsent((authenticationInfo
+                                                    .getTenantDomain() + "/" + deviceType + "/" + deviceId),
+                                            clientRemote);
+                            if (activeSession != null && activeSession.getMySession().isOpen() && activeSession
+                                    .getPeerSession() == null) {
+                                throw new RemoteSessionManagementException("Another client session waiting on device " +
+                                        "to connect.");
                             } else {
-                                Session lastSession = RemoteSessionManagementDataHolder.getInstance().getDeviceRequestMap().putIfAbsent(
-                                        (authenticationInfo.getTenantDomain() + "/" + deviceType + "/" + deviceId),
-                                        session);
+                                // if there is pending session exists but already closed, then we need to remove it.
+                                if (activeSession != null) {
+                                    endSession(activeSession.getMySession(), "Remote session closed due to new session" +
+                                            " request");
+                                    // Use put if absent for adding session to waiting list since we need to overcome
+                                    // multithreaded session requests.
+                                    activeSession = RemoteSessionManagementDataHolder.getInstance()
+                                            .getActiveDeviceClientSessionMap().putIfAbsent((authenticationInfo
+                                                    .getTenantDomain() + "/" + deviceType + "/" +
+                                                    deviceId), clientRemote);
+                                }
 
-                                if (lastSession == null) {
+                                // If another client tried to start session same time then active session will be
+                                // exist. So we are adding session request only no parallel sessions added to map
+                                if (activeSession == null) {
 
                                     // Create operation if session initiated by client
                                     Operation operation = new ConfigOperation();
@@ -164,55 +171,58 @@ public class RemoteSessionManagementServiceImpl implements RemoteSessionManageme
                                     operation.setEnabled(true);
                                     operation.setControl(Operation.Control.NO_REPEAT);
                                     JSONObject payload = new JSONObject();
-                                    payload.put("serverUrl", RemoteSessionManagementDataHolder.getInstance().getServerUrl());
+                                    payload.put("serverUrl", RemoteSessionManagementDataHolder.getInstance()
+                                            .getServerUrl());
                                     operation.setPayLoad(payload.toString());
-                                    String date = new SimpleDateFormat(DATE_FORMAT_NOW).format(new Date());
+                                    String date = new SimpleDateFormat(RemoteSessionConstants.DATE_FORMAT_NOW).format
+                                            (new Date());
                                     operation.setCreatedTimeStamp(date);
-
                                     List<DeviceIdentifier> deviceIdentifiers = new ArrayList<>();
                                     deviceIdentifiers.add(new DeviceIdentifier(deviceId, deviceType));
                                     Activity activity = RemoteSessionManagementDataHolder.getInstance()
                                             .getDeviceManagementProviderService().addOperation(deviceType, operation,
                                                     deviceIdentifiers);
-                                    log.info("Activity id: " + activity.getActivityId());
+                                    clientRemote.setOperationId(activity.getActivityId()
+                                            .replace(DeviceManagementConstants.OperationAttributes.ACTIVITY, ""));
+                                    RemoteSessionManagementDataHolder.getInstance().getSessionMap().put(session.getId
+                                            (), clientRemote);
+                                    log.info("Client remote session opened for session id: " + session.getId() +
+                                            " device Type : " + deviceType + " , " + "deviceId : " + deviceId);
 
-                                    RemoteSession clientRemote = new ClientSession(session, authenticationInfo
-                                            .getTenantDomain(), deviceType, deviceId, activity.getActivityId().replace(DeviceManagementConstants
-                                            .OperationAttributes.ACTIVITY, ""));
-                                    RemoteSessionManagementDataHolder.getInstance().getSessionMap().put(session.getId(), clientRemote);
+                                } else {
+                                    throw new RemoteSessionManagementException("Another client session waiting on " +
+                                            "device to connect.");
                                 }
                             }
                         }
-                        log.info("Current session count: " + RemoteSessionManagementDataHolder
-                                .getInstance().getSessionMap().size());
+                        log.info("Current remote sessions count: " + RemoteSessionManagementDataHolder.getInstance()
+                                .getSessionMap().size());
 
                     } else {
-                        throw new RemoteSessionInvalidException("Missing device Id or type ", new CloseReason
-                                (CloseReason.CloseCodes.CANNOT_ACCEPT, "Missing device Id or device type "));
+                        throw new RemoteSessionManagementException("Missing device Id or type ");
                     }
                 } else {
-                    throw new RemoteSessionInvalidException("Unauthorized Access for the device Type : " + deviceType
-                            + " , deviceId : " + deviceId, new CloseReason(CloseReason.CloseCodes.CANNOT_ACCEPT,
-                            "Unauthorized Access"));
+                    throw new RemoteSessionManagementException("Unauthorized Access for the device Type : " + deviceType
+                            + " , deviceId : " + deviceId);
                 }
             } catch (OperationManagementException | InvalidDeviceException e) {
-                throw new RemoteSessionManagementException("Error occurred while adding initial operation for the device Type : " +
-                        deviceType + " , deviceId : " + deviceId, e);
+                throw new RemoteSessionManagementException("Error occurred while adding initial operation for the " +
+                        "device Type : " + deviceType + " , deviceId : " + deviceId);
             } catch (DeviceAccessAuthorizationException e) {
-                throw new RemoteSessionManagementException("Error occurred while device access authorization for the device Type : " +
-                        deviceType + " , " + "deviceId : " + deviceId, e);
+                throw new RemoteSessionManagementException("Error occurred while device access authorization for the " +
+                        "device Type : " + deviceType + " , " + "deviceId : " + deviceId);
             } finally {
                 PrivilegedCarbonContext.endTenantFlow();
             }
 
         } else {
-            throw new RemoteSessionInvalidException("Invalid token", new CloseReason(CloseReason.CloseCodes
-                    .CANNOT_ACCEPT, "Invalid token"));
+            throw new RemoteSessionManagementException("Invalid token");
         }
     }
 
     @Override
-    public void initializeSession(Session session, String deviceType, String deviceId) throws RemoteSessionInvalidException, RemoteSessionManagementException {
+    public void initializeSession(Session session, String deviceType, String deviceId) throws
+            RemoteSessionManagementException {
         initializeSession(session, deviceType, deviceId, null);
     }
 
@@ -221,22 +231,18 @@ public class RemoteSessionManagementServiceImpl implements RemoteSessionManageme
      *
      * @param session Web socket RemoteSession
      * @param message String message needs to send to peer connection
-     * @throws RemoteSessionInvalidException    throws when session cannot be made due to invalid data
+     * @throws RemoteSessionManagementException throws when session cannot be made due to invalid data
      * @throws RemoteSessionManagementException throws when session has error with accessing device resources
      */
     @Override
-    public void sendMessageToPeer(Session session, String message) throws RemoteSessionManagementException,
-            RemoteSessionInvalidException {
+    public void sendMessageToPeer(Session session, String message) throws RemoteSessionManagementException {
         JSONObject jsonObject = new JSONObject(message);
-        RemoteSession remoteSession = RemoteSessionManagementDataHolder.getInstance().getSessionMap().get(session.getId());
+        RemoteSession remoteSession = RemoteSessionManagementDataHolder.getInstance().getSessionMap().get(session
+                .getId());
         if (remoteSession != null) {
-            if (remoteSession instanceof ClientSession) {
-                jsonObject.put("id", remoteSession.getOperationId());
-            }
             remoteSession.sendMessageToPeer(jsonObject.toString());
         } else {
-            throw new RemoteSessionInvalidException("Remote Session cannot be found ", new CloseReason(CloseReason
-                    .CloseCodes.CANNOT_ACCEPT, "Invalid RemoteSession"));
+            throw new RemoteSessionManagementException("Remote Session cannot be found ");
         }
     }
 
@@ -246,48 +252,68 @@ public class RemoteSessionManagementServiceImpl implements RemoteSessionManageme
      *
      * @param session Web socket RemoteSession
      * @param message Byte message needs to send to peer connection
-     * @throws RemoteSessionInvalidException    throws when session cannot be made due to invalid data
+     * @throws RemoteSessionManagementException throws when session cannot be made due to invalid data
      * @throws RemoteSessionManagementException throws when session has error with accessing device resources
      */
     @Override
-    public void sendMessageToPeer(Session session, byte[] message) throws RemoteSessionInvalidException,
-            RemoteSessionManagementException {
+    public void sendMessageToPeer(Session session, byte[] message) throws RemoteSessionManagementException {
 
-        RemoteSession remoteSession = RemoteSessionManagementDataHolder.getInstance().getSessionMap().get(session.getId());
+        RemoteSession remoteSession = RemoteSessionManagementDataHolder.getInstance().getSessionMap().get(session
+                .getId());
         if (remoteSession != null) {
             remoteSession.sendMessageToPeer(message);
         } else {
-            throw new RemoteSessionInvalidException("Remote Session cannot be found ", new CloseReason(CloseReason
-                    .CloseCodes.CANNOT_ACCEPT, "Invalid RemoteSession"));
+            throw new RemoteSessionManagementException("Remote Session cannot be found ");
         }
     }
 
     /**
      * Closing the session and cleanup the resources
      *
-     * @param session Web socket RemoteSession
+     * @param session Web socket Remote Session
      */
     @Override
-    public void endSession(Session session) throws IOException {
-
-        RemoteSession remoteSession = RemoteSessionManagementDataHolder.getInstance().getSessionMap().remove(session.getId());
+    public void endSession(Session session, String closeReason) {
+        log.info("Closing session: "+session.getId()+" due to:"+ closeReason);
+        RemoteSession remoteSession = RemoteSessionManagementDataHolder.getInstance().getSessionMap().remove(session
+                .getId());
         if (remoteSession != null) {
             String operationId = remoteSession.getOperationId();
-            Session peerSession = remoteSession.getPeerSession().getMySession();
-            if (peerSession != null) {
-                RemoteSessionManagementDataHolder.getInstance().getSessionMap().remove(peerSession.getId());
-                if (peerSession.isOpen()) {
-                    peerSession.close(new CloseReason(CloseReason.CloseCodes.GOING_AWAY, "Remote session closed"));
+            if (remoteSession.getPeerSession() != null) {
+                Session peerSession = remoteSession.getPeerSession().getMySession();
+                if (peerSession != null) {
+                    RemoteSessionManagementDataHolder.getInstance().getSessionMap().remove(peerSession.getId());
+                    if (peerSession.isOpen()) {
+                        try {
+                            peerSession.close(new CloseReason(CloseReason.CloseCodes.GOING_AWAY, closeReason));
+                        } catch (IOException ex) {
+                            if (log.isDebugEnabled()) {
+                                log.error("Failed to disconnect the client.", ex);
+                            }
+                        }
+                    }
+                }
+            }
+            if (remoteSession.getMySession() != null) {
+                Session mySession = remoteSession.getMySession();
+                if (mySession.isOpen()) {
+                    try {
+                        mySession.close(new CloseReason(CloseReason.CloseCodes.GOING_AWAY, closeReason));
+                    } catch (IOException ex) {
+                        if (log.isDebugEnabled()) {
+                            log.error("Failed to disconnect the client.", ex);
+                        }
+                    }
                 }
             }
             if (operationId != null) {
-                Session lastSession = RemoteSessionManagementDataHolder.getInstance().getDeviceRequestMap().get(
-                        (remoteSession.getTenantDomain() + "/" + remoteSession.getDeviceType() + "/" + remoteSession
-                                .getDeviceId()));
-                if (lastSession != null && lastSession.getId().equals(session.getId())) {
-                    RemoteSessionManagementDataHolder.getInstance().getDeviceRequestMap().remove(
-                            (remoteSession.getTenantDomain() + "/" + remoteSession.getDeviceType() + "/" + remoteSession
-                                    .getDeviceId()));
+                String deviceIdentifier = remoteSession.getTenantDomain() + "/" + remoteSession
+                        .getDeviceType() + "/" + remoteSession.getDeviceId();
+                RemoteSession lastSession = RemoteSessionManagementDataHolder.getInstance()
+                        .getActiveDeviceClientSessionMap().get(deviceIdentifier);
+                if (lastSession != null && lastSession.getMySession().getId().equals(session.getId())) {
+                    RemoteSessionManagementDataHolder.getInstance().getActiveDeviceClientSessionMap().remove
+                            (deviceIdentifier);
                 }
             }
         }
diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/authentication/AuthenticationInfo.java b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/authentication/AuthenticationInfo.java
index 9da426df3c..e7d77e5bdb 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/authentication/AuthenticationInfo.java
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/authentication/AuthenticationInfo.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
+ * Copyright (c) 2017, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/authentication/OAuthAuthenticator.java b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/authentication/OAuthAuthenticator.java
index 689dd159c9..b6832fc488 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/authentication/OAuthAuthenticator.java
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/authentication/OAuthAuthenticator.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
+ * Copyright (c) 2017, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/authentication/oauth/OAuthTokenValidator.java b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/authentication/oauth/OAuthTokenValidator.java
index ab19f5e08e..3eb6d0db28 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/authentication/oauth/OAuthTokenValidator.java
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/authentication/oauth/OAuthTokenValidator.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
+ * Copyright (c) 2017, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/authentication/oauth/OAuthTokenValidatorStubFactory.java b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/authentication/oauth/OAuthTokenValidatorStubFactory.java
index 8b02b920c3..9ee5a76936 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/authentication/oauth/OAuthTokenValidatorStubFactory.java
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/authentication/oauth/OAuthTokenValidatorStubFactory.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
+ * Copyright (c) 2017, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
  *
  * WSO2 Inc. licenses this file to you under the Apache License,
  * Version 2.0 (the "License"); you may not use this file except
@@ -54,7 +54,6 @@ public class OAuthTokenValidatorStubFactory extends BasePoolableObjectFactory {
     private HttpClient httpClient;
     Map<String, String> tokenValidationProperties;
 
-
     public OAuthTokenValidatorStubFactory(Map<String, String> globalProperties) {
         this.tokenValidationProperties = globalProperties;
         this.httpClient = createHttpClient();
diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/authentication/oauth/exception/OAuthTokenValidationException.java b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/authentication/oauth/exception/OAuthTokenValidationException.java
index 3e57154766..069aa11876 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/authentication/oauth/exception/OAuthTokenValidationException.java
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/authentication/oauth/exception/OAuthTokenValidationException.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
+ * Copyright (c) 2017, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
  *
  * WSO2 Inc. licenses this file to you under the Apache License,
  * Version 2.0 (the "License"); you may not use this file except
diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/constants/RemoteSessionConstants.java b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/constants/RemoteSessionConstants.java
index 15f9141d52..07ebcd5082 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/constants/RemoteSessionConstants.java
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/constants/RemoteSessionConstants.java
@@ -1,5 +1,5 @@
 /*
- *  Copyright (c) 2016, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
+ *  Copyright (c) 2017, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
  *
  *  WSO2 Inc. licenses this file to you under the Apache License,
  *  Version 2.0 (the "License"); you may not use this file except
@@ -19,7 +19,7 @@
 package org.wso2.carbon.device.mgt.extensions.remote.session.constants;
 
 /**
- * This holds the constants related to this feature
+ * This holds the constants related to remote session
  */
 public class RemoteSessionConstants {
 
@@ -28,14 +28,18 @@ public class RemoteSessionConstants {
     public static final String TOKEN_VALIDATION_CONTEX = "/services/OAuth2TokenValidationService";
     public static final String USERNAME = "username";
     public static final String PASSWORD = "password";
-    public static final String REMOTE_CONNECT = "REMOTE_CONNECT";
-    public static final String THROTTLE_OUT = "THROTTLE_OUT";
+    public static final String DATE_FORMAT_NOW = "yyyy-MM-dd HH:mm:ss";
+    public static final int MAX_BUFFER_SIZE = 640 * 1024;
+
     public static final String QUERY_STRING = "queryString";
     public static final String MAXIMUM_TOTAL_HTTP_CONNECTION = "maximumTotalHttpConnection";
     public static final String MAXIMUM_HTTP_CONNECTION_PER_HOST = "maximumHttpConnectionPerHost";
     public static final String DEFAULT_MAXIMUM_HTTP_CONNECTION_PER_HOST = "2";
     public static final String DEFAULT_MAXIMUM_TOTAL_HTTP_CONNECTIONS = "100";
 
+    // Initial remote connection operation related Constants
+    public static final String REMOTE_CONNECT = "REMOTE_CONNECT";
+    public static final String REMOTE_CONNECT_CODE = "code";
 
     private RemoteSessionConstants() {
     }
diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/dto/RemoteSession.java b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/dto/RemoteSession.java
index 5402f131f2..fe92926193 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/dto/RemoteSession.java
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/dto/RemoteSession.java
@@ -15,41 +15,44 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-package org.wso2.carbon.device.mgt.extensions.remote.session.dto.common;
+package org.wso2.carbon.device.mgt.extensions.remote.session.dto;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
-import org.wso2.carbon.device.mgt.extensions.remote.session.exception.RemoteSessionInvalidException;
 import org.wso2.carbon.device.mgt.extensions.remote.session.exception.RemoteSessionManagementException;
+import org.wso2.carbon.device.mgt.extensions.remote.session.internal.RemoteSessionManagementDataHolder;
 
-import javax.websocket.CloseReason;
 import javax.websocket.Session;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 
 /**
  * {@link RemoteSession} will represent remote websocket session
- * This class implements the behaviours of sending message to the session in multithreaded environment.
+ * This class implements the behaviours of sending message to the session in multithreaded context.
  */
 public class RemoteSession {
 
     private static final Log log = LogFactory.getLog(RemoteSession.class);
     private String tenantDomain, operationId, deviceType, deviceId;
     private long lastMessageTimeStamp = System.currentTimeMillis();
-    ;
     private RemoteSession peerSession;
     private Session mySession;
     private final Object writeLockObject = new Object();
+    private int maxMessagesPerSecond;
+    private int messageAllowance;
+    private double messageRatePerSecond;
 
-    protected RemoteSession(Session session, String tenantDomain, String deviceType, String deviceId) {
+    public RemoteSession(Session session, String tenantDomain, String deviceType, String deviceId) {
         this.mySession = session;
         this.deviceType = deviceType;
         this.deviceId = deviceId;
         this.tenantDomain = tenantDomain;
+        maxMessagesPerSecond = RemoteSessionManagementDataHolder.getInstance().getMaxMessagesPerSecond();
+        messageAllowance = maxMessagesPerSecond;
+        messageRatePerSecond = (double) maxMessagesPerSecond / 1000;
     }
 
-    private void sendMessage(Object message) throws RemoteSessionInvalidException, RemoteSessionManagementException {
-
+    private void sendMessage(Object message) throws RemoteSessionManagementException {
         if (message != null) {
             boolean isMessageCountExceed = false;
             if (mySession != null && mySession.isOpen()) {
@@ -71,38 +74,41 @@ public class RemoteSession {
                     }
                 }
             } else {
-                throw new RemoteSessionInvalidException("Peer Session already closed ", new CloseReason
-                        (CloseReason.CloseCodes.CANNOT_ACCEPT, "Peer Session already closed "));
+                throw new RemoteSessionManagementException("Peer Session already closed ");
             }
         } else {
             throw new RemoteSessionManagementException("Message is empty");
         }
     }
 
-    public void sendMessageToPeer(Object message) throws RemoteSessionInvalidException,
-            RemoteSessionManagementException {
+    public void sendMessageToPeer(Object message) throws RemoteSessionManagementException {
         peerSession.sendMessage(message);
     }
 
-
     /**
      * Use for limit the messages for given time
      *
      * @return message rate applied
      */
-    public boolean applyRateLimit(){
-        return false;
+    private boolean applyRateLimit() {
+        long currentTime = System.currentTimeMillis();
+        messageAllowance += (currentTime - lastMessageTimeStamp) * messageRatePerSecond;
+        if (messageAllowance > maxMessagesPerSecond) {
+            messageAllowance = maxMessagesPerSecond;
+        }
+        if (messageAllowance >= 1) {
+            lastMessageTimeStamp = currentTime;
+            messageAllowance -= 1;
+            return false;
+        } else {
+            return true;
+        }
     }
 
-
     public Session getMySession() {
         return mySession;
     }
 
-    public void setMySession(Session mySession) {
-        this.mySession = mySession;
-    }
-
     public RemoteSession getPeerSession() {
         return peerSession;
     }
@@ -115,10 +121,6 @@ public class RemoteSession {
         return tenantDomain;
     }
 
-    public void setTenantDomain(String tenantDomain) {
-        this.tenantDomain = tenantDomain;
-    }
-
     public String getOperationId() {
         return operationId;
     }
@@ -131,23 +133,7 @@ public class RemoteSession {
         return deviceType;
     }
 
-    public void setDeviceType(String deviceType) {
-        this.deviceType = deviceType;
-    }
-
     public String getDeviceId() {
         return deviceId;
     }
-
-    public void setDeviceId(String deviceId) {
-        this.deviceId = deviceId;
-    }
-
-    public long getLastMessageTimeStamp() {
-        return lastMessageTimeStamp;
-    }
-
-    public void setLastMessageTimeStamp(long lastMessageTimeStamp) {
-        this.lastMessageTimeStamp = lastMessageTimeStamp;
-    }
 }
diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/exception/RemoteSessionManagementException.java b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/exception/RemoteSessionManagementException.java
index 628fed35ba..741f04e31a 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/exception/RemoteSessionManagementException.java
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/exception/RemoteSessionManagementException.java
@@ -18,28 +18,23 @@
 package org.wso2.carbon.device.mgt.extensions.remote.session.exception;
 
 
+import javax.websocket.CloseReason;
+
 /**
- * This Exception will be thrown, when there any interference with Remote RemoteSession.
+ * This Exception will be thrown, when there any management issues with Remote Session.
  */
 public class RemoteSessionManagementException extends Exception {
 
-    public RemoteSessionManagementException(String msg, Exception nestedEx) {
-        super(msg, nestedEx);
-    }
-
-    public RemoteSessionManagementException(String message, Throwable cause) {
-        super(message, cause);
-    }
+    CloseReason closeReason;
 
     public RemoteSessionManagementException(String msg) {
         super(msg);
+        this.closeReason = new CloseReason(CloseReason.CloseCodes.CANNOT_ACCEPT, msg);
     }
 
-    public RemoteSessionManagementException() {
-        super();
+    public CloseReason getCloseReason() {
+        return closeReason;
     }
 
-    public RemoteSessionManagementException(Throwable cause) {
-        super(cause);
-    }
+
 }
diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/internal/RemoteSessionManagementDataHolder.java b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/internal/RemoteSessionManagementDataHolder.java
index c015e50e0a..310affbb0c 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/internal/RemoteSessionManagementDataHolder.java
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/internal/RemoteSessionManagementDataHolder.java
@@ -21,12 +21,15 @@ package org.wso2.carbon.device.mgt.extensions.remote.session.internal;
 import org.wso2.carbon.device.mgt.common.authorization.DeviceAccessAuthorizationService;
 import org.wso2.carbon.device.mgt.core.service.DeviceManagementProviderService;
 import org.wso2.carbon.device.mgt.extensions.remote.session.authentication.OAuthAuthenticator;
-import org.wso2.carbon.device.mgt.extensions.remote.session.dto.common.RemoteSession;
+import org.wso2.carbon.device.mgt.extensions.remote.session.dto.RemoteSession;
 
-import javax.websocket.Session;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
+/**
+ * Class {@link RemoteSessionManagementDataHolder} will hold the configurations and in memory storage strictures to
+ * manage remote sessions.
+ */
 public class RemoteSessionManagementDataHolder {
 
     private static RemoteSessionManagementDataHolder thisInstance = new RemoteSessionManagementDataHolder();
@@ -35,9 +38,9 @@ public class RemoteSessionManagementDataHolder {
     private boolean isEnabled;
     private String serverUrl;
     private long maxIdleTimeout;
-    private int messagesPerSession;
+    private int maxMessagesPerSecond;
     private OAuthAuthenticator oAuthAuthenticator;
-    private Map<String, Session> deviceRequestMap = new ConcurrentHashMap<String, Session>();
+    private Map<String, RemoteSession> activeDeviceClientSessionMap = new ConcurrentHashMap<String, RemoteSession>();
     private Map<String, RemoteSession> sessionMap = new ConcurrentHashMap<String, RemoteSession>();
 
     public static RemoteSessionManagementDataHolder getInstance() {
@@ -72,16 +75,8 @@ public class RemoteSessionManagementDataHolder {
         return sessionMap;
     }
 
-    public void setSessionMap(Map<String, RemoteSession> sessionMap) {
-        this.sessionMap = sessionMap;
-    }
-
-    public Map<String, Session> getDeviceRequestMap() {
-        return deviceRequestMap;
-    }
-
-    public void setDeviceRequestMap(Map<String, Session> deviceRequestMap) {
-        this.deviceRequestMap = deviceRequestMap;
+    public Map<String, RemoteSession> getActiveDeviceClientSessionMap() {
+        return activeDeviceClientSessionMap;
     }
 
     public boolean isEnabled() {
@@ -100,12 +95,12 @@ public class RemoteSessionManagementDataHolder {
         this.serverUrl = serverUrl;
     }
 
-    public int getMessagesPerSession() {
-        return messagesPerSession;
+    public int getMaxMessagesPerSecond() {
+        return maxMessagesPerSecond;
     }
 
-    public void setMessagesPerSession(int messagesPerSession) {
-        this.messagesPerSession = messagesPerSession;
+    public void setMaxMessagesPerSecond(int maxMessagesPerSecond) {
+        this.maxMessagesPerSecond = maxMessagesPerSecond;
     }
 
     public long getMaxIdleTimeout() {
diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/internal/RemoteSessionManagementServiceComponent.java b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/internal/RemoteSessionManagementServiceComponent.java
index 1bdb3e17ab..2d3e913fc7 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/internal/RemoteSessionManagementServiceComponent.java
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/internal/RemoteSessionManagementServiceComponent.java
@@ -54,10 +54,10 @@ public class RemoteSessionManagementServiceComponent {
         try {
 
             BundleContext bundleContext = componentContext.getBundleContext();
-            bundleContext.registerService(ServerStartupObserver.class.getName(), new RemoteSessionManagerStartupListener(),
-                    null);
-            bundleContext.registerService(RemoteSessionManagementService.class.getName(), new RemoteSessionManagementServiceImpl(),
-                    null);
+            bundleContext.registerService(ServerStartupObserver.class.getName(), new
+                    RemoteSessionManagerStartupListener(), null);
+            bundleContext.registerService(RemoteSessionManagementService.class.getName(), new
+                    RemoteSessionManagementServiceImpl(), null);
             if (log.isDebugEnabled()) {
                 log.debug("Remote Session device access service implementation bundle has been successfully " +
                         "initialized");
@@ -73,18 +73,23 @@ public class RemoteSessionManagementServiceComponent {
     }
 
     protected void setDeviceManagementProviderService(DeviceManagementProviderService deviceManagementProviderService) {
-        RemoteSessionManagementDataHolder.getInstance().setDeviceManagementProviderService(deviceManagementProviderService);
+        RemoteSessionManagementDataHolder.getInstance()
+                .setDeviceManagementProviderService(deviceManagementProviderService);
     }
 
-    protected void unsetDeviceManagementProviderService(DeviceManagementProviderService deviceManagementProviderService) {
+    protected void unsetDeviceManagementProviderService(DeviceManagementProviderService
+                                                                deviceManagementProviderService) {
         RemoteSessionManagementDataHolder.getInstance().setDeviceManagementProviderService(null);
     }
 
-    protected void setDeviceAccessAuthorizationService(DeviceAccessAuthorizationService deviceAccessAuthorizationService) {
-        RemoteSessionManagementDataHolder.getInstance().setDeviceAccessAuthorizationService(deviceAccessAuthorizationService);
+    protected void setDeviceAccessAuthorizationService(DeviceAccessAuthorizationService
+                                                               deviceAccessAuthorizationService) {
+        RemoteSessionManagementDataHolder.getInstance()
+                .setDeviceAccessAuthorizationService(deviceAccessAuthorizationService);
     }
 
-    protected void unsetDeviceAccessAuthorizationService(DeviceAccessAuthorizationService deviceAccessAuthorizationService) {
+    protected void unsetDeviceAccessAuthorizationService(DeviceAccessAuthorizationService
+                                                                 deviceAccessAuthorizationService) {
         RemoteSessionManagementDataHolder.getInstance().setDeviceManagementProviderService(null);
     }
 
diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/listener/RemoteSessionManagerStartupListener.java b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/listener/RemoteSessionManagerStartupListener.java
index a02b945ee4..19b89f369c 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/listener/RemoteSessionManagerStartupListener.java
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/listener/RemoteSessionManagerStartupListener.java
@@ -62,10 +62,14 @@ public class RemoteSessionManagerStartupListener implements ServerStartupObserve
             Map<String, String> configProperties = new HashMap<>();
 
             // Set max idle timeout in milliseconds
-            RemoteSessionManagementDataHolder.getInstance().setMaxIdleTimeout(rsConfig.getSessionIdleTimeOut()*60000);
+            RemoteSessionManagementDataHolder.getInstance().setMaxIdleTimeout((long) rsConfig.getSessionIdleTimeOut() *
+                    60000);
 
             // Set max messages per second.
-            RemoteSessionManagementDataHolder.getInstance().setMessagesPerSession(rsConfig.getMaxMessagesPerSession());
+            if (rsConfig.getMaxMessagesPerSession() > 0) {
+                RemoteSessionManagementDataHolder.getInstance().setMaxMessagesPerSecond(rsConfig
+                        .getMaxMessagesPerSession());
+            }
             // Token validation related configuration
             configProperties.put(RemoteSessionConstants.TOKEN_VALIDATION_ENDPOINT_URL, kmConfig.getServerUrl());
             configProperties.put(RemoteSessionConstants.USERNAME, kmConfig.getAdminUsername());
diff --git a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/util/PropertyUtils.java b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/util/PropertyUtils.java
index b8e46ad41e..cb63b97ac8 100644
--- a/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/util/PropertyUtils.java
+++ b/components/extensions/remote-session-extension/org.wso2.carbon.device.mgt.extensions.remote.session/src/main/java/org/wso2/carbon/device/mgt/extensions/remote.session/util/PropertyUtils.java
@@ -1,5 +1,5 @@
 /*
-*  Copyright (c) 2016, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
+*  Copyright (c) 2017, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
 *
 *  WSO2 Inc. licenses this file to you under the Apache License,
 *  Version 2.0 (the "License"); you may not use this file except
@@ -23,9 +23,17 @@ import org.wso2.carbon.device.mgt.extensions.remote.session.exception.RemoteSess
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
+/**
+ * Class for reading web socket url parameters
+ */
 public class PropertyUtils {
 
-    //This method is only used if the mb features are within DAS.
+    /**
+     * Replace URL with placeholders with properties
+     * @param urlWithPlaceholders URL
+     * @return replaced url
+     * @throws RemoteSessionManagementException
+     */
     public static String replaceProperty(String urlWithPlaceholders) throws RemoteSessionManagementException {
         String regex = "\\$\\{(.*?)\\}";
         Pattern pattern = Pattern.compile(regex);
diff --git a/components/extensions/remote-session-extension/pom.xml b/components/extensions/remote-session-extension/pom.xml
index 42b540708b..44d6ca0876 100644
--- a/components/extensions/remote-session-extension/pom.xml
+++ b/components/extensions/remote-session-extension/pom.xml
@@ -22,7 +22,7 @@
     <parent>
         <groupId>org.wso2.carbon.devicemgt-plugins</groupId>
         <artifactId>extensions</artifactId>
-        <version>4.0.54-SNAPSHOT</version>
+        <version>4.0.66-SNAPSHOT</version>
         <relativePath>../pom.xml</relativePath>
     </parent>