1 files changed, 14 insertions, 6 deletions
diff --git a/lib/realtime.js b/lib/realtime.js
index c731e5b0..d6ba62b2 100644
--- a/lib/realtime.js
+++ b/lib/realtime.js
@@ -272,16 +272,24 @@ function isReady () {
 }
 
 function extractNoteIdFromSocket (socket) {
-  if (!socket || !socket.handshake || !socket.handshake.headers) {
+  if (!socket || !socket.handshake) {
     return false
   }
-  var referer = socket.handshake.headers.referer
-  if (!referer) {
+  if (socket.handshake.query && socket.handshake.query.noteId) {
+    return socket.handshake.query.noteId
+  } else if (socket.handshake.headers) {
+    // this part is only for backward compatibility only; current code
+    // should be using noteId query parameter instead.
+    var referer = socket.handshake.headers.referer
+    if (!referer) {
+      return false
+    }
+    var hostUrl = url.parse(referer)
+    var noteId = config.urlpath ? hostUrl.pathname.slice(config.urlpath.length + 1, hostUrl.pathname.length).split('/')[1] : hostUrl.pathname.split('/')[1]
+    return noteId
+  } else {
     return false
   }
-  var hostUrl = url.parse(referer)
-  var noteId = config.urlpath ? hostUrl.pathname.slice(config.urlpath.length + 1, hostUrl.pathname.length).split('/')[1] : hostUrl.pathname.split('/')[1]
-  return noteId
 }
 
 function parseNoteIdFromSocket (socket, callback) {