From 8b10c68d9fc9ce8f86eda5a31273b22960bbafa0 2014-10-09 22:12:20
From: MinRK <benjaminrk@gmail.com>
Date: 2014-10-09 22:12:20
Subject: [PATCH] remove on_first_message authentication

in ZMQStreams. No need for weird, special first message.

- use regular cookie auth
- use url param for session id

---

diff --git a/IPython/html/base/zmqhandlers.py b/IPython/html/base/zmqhandlers.py
index f9f1eb1..b0caf56 100644
--- a/IPython/html/base/zmqhandlers.py
+++ b/IPython/html/base/zmqhandlers.py
@@ -146,18 +146,31 @@ class AuthenticatedZMQStreamHandler(ZMQStreamHandler, IPythonHandler):
         which doesn't make sense for websockets
         """
         pass
-
-    def open(self, kernel_id):
-        self.kernel_id = cast_unicode(kernel_id, 'ascii')
+    
+    def get(self, *args, **kwargs):
         # Check to see that origin matches host directly, including ports
         # Tornado 4 already does CORS checking
         if tornado.version_info[0] < 4:
             if not self.check_origin(self.get_origin()):
                 raise web.HTTPError(403)
-
+        
+        # authenticate the request before opening the websocket
+        if self.get_current_user() is None:
+            self.log.warn("Couldn't authenticate WebSocket connection")
+            raise web.HTTPError(403)
+        
+        if self.get_argument('session_id'):
+            self.session.session = cast_unicode(self.get_argument('session_id'))
+        else:
+            self.log.warn("No session ID specified")
+        
+        return super(AuthenticatedZMQStreamHandler, self).get(*args, **kwargs)
+    
+    def initialize(self):
         self.session = Session(config=self.config)
-        self.save_on_message = self.on_message
-        self.on_message = self.on_first_message
+    
+    def open(self, kernel_id):
+        self.kernel_id = cast_unicode(kernel_id, 'ascii')
         
         # start the pinging
         if self.ping_interval > 0:
@@ -187,28 +200,3 @@ class AuthenticatedZMQStreamHandler(ZMQStreamHandler, IPythonHandler):
 
     def on_pong(self, data):
         self.last_pong = ioloop.IOLoop.instance().time()
-
-    def _inject_cookie_message(self, msg):
-        """Inject the first message, which is the document cookie,
-        for authentication."""
-        if not PY3 and isinstance(msg, unicode):
-            # Cookie constructor doesn't accept unicode strings
-            # under Python 2.x for some reason
-            msg = msg.encode('utf8', 'replace')
-        try:
-            identity, msg = msg.split(':', 1)
-            self.session.session = cast_unicode(identity, 'ascii')
-        except Exception:
-            logging.error("First ws message didn't have the form 'identity:[cookie]' - %r", msg)
-        
-        try:
-            self.request._cookies = SimpleCookie(msg)
-        except:
-            self.log.warn("couldn't parse cookie string: %s",msg, exc_info=True)
-
-    def on_first_message(self, msg):
-        self._inject_cookie_message(msg)
-        if self.get_current_user() is None:
-            self.log.warn("Couldn't authenticate WebSocket connection")
-            raise web.HTTPError(403)
-        self.on_message = self.save_on_message
diff --git a/IPython/html/services/kernels/handlers.py b/IPython/html/services/kernels/handlers.py
index b51861e..9406796 100644
--- a/IPython/html/services/kernels/handlers.py
+++ b/IPython/html/services/kernels/handlers.py
@@ -126,16 +126,12 @@ class ZMQChannelHandler(AuthenticatedZMQStreamHandler):
         self.kernel_info_channel.close()
         self.kernel_info_channel = None
     
-    
-    def initialize(self, *args, **kwargs):
+    def initialize(self):
+        super(ZMQChannelHandler, self).initialize()
         self.zmq_stream = None
     
-    def on_first_message(self, msg):
-        try:
-            super(ZMQChannelHandler, self).on_first_message(msg)
-        except web.HTTPError:
-            self.close()
-            return
+    def open(self, kernel_id):
+        super(ZMQChannelHandler, self).open(kernel_id)
         try:
             self.create_stream()
         except web.HTTPError:
diff --git a/IPython/html/static/services/kernels/js/kernel.js b/IPython/html/static/services/kernels/js/kernel.js
index 6c9187c..3a25e90 100644
--- a/IPython/html/static/services/kernels/js/kernel.js
+++ b/IPython/html/static/services/kernels/js/kernel.js
@@ -424,16 +424,17 @@ define([
         var ws_host_url = this.ws_url + this.kernel_url;
 
         console.log("Starting WebSockets:", ws_host_url);
-
-        this.channels.shell = new this.WebSocket(
-            this.ws_url + utils.url_join_encode(this.kernel_url, "shell")
-        );
-        this.channels.stdin = new this.WebSocket(
-            this.ws_url + utils.url_join_encode(this.kernel_url, "stdin")
-        );
-        this.channels.iopub = new this.WebSocket(
-            this.ws_url + utils.url_join_encode(this.kernel_url, "iopub")
-        );
+        
+        var channel_url = function(channel) {
+            return [
+                that.ws_url,
+                utils.url_join_encode(that.kernel_url, channel),
+                "?session_id=" + that.session_id
+            ].join('');
+        };
+        this.channels.shell = new this.WebSocket(channel_url("shell"));
+        this.channels.stdin = new this.WebSocket(channel_url("stdin"));
+        this.channels.iopub = new this.WebSocket(channel_url("iopub"));
         
         var already_called_onclose = false; // only alert once
         var ws_closed_early = function(evt){
@@ -492,16 +493,12 @@ define([
     };
 
     /**
-     * Handle a websocket entering the open state sends session and
-     * cookie authentication info as first message.
+     * Handle a websocket entering the open state,
+     * signaling that the kernel is connected when all channels are open.
      *
      * @function _ws_opened
      */
     Kernel.prototype._ws_opened = function (evt) {
-        // send the session id so the Session object Python-side
-        // has the same identity
-        evt.target.send(this.session_id + ':' + document.cookie);
-
         if (this.is_connected()) {
             // all events ready, trigger started event.
             this._kernel_connected();