authenticate Websockets with the session cookie

Now all Notebook connections are authenticated.
This commit is contained in:
MinRK 2011-08-30 16:27:12 -07:00
parent 152fe1b6f1
commit d11ef986a4
2 changed files with 65 additions and 9 deletions

View File

@ -16,6 +16,9 @@ Authors:
# Imports # Imports
#----------------------------------------------------------------------------- #-----------------------------------------------------------------------------
import logging
import Cookie
from tornado import web from tornado import web
from tornado import websocket from tornado import websocket
@ -165,20 +168,58 @@ class ZMQStreamHandler(websocket.WebSocketHandler):
else: else:
self.write_message(msg) self.write_message(msg)
class AuthenticatedZMQStreamHandler(ZMQStreamHandler):
def open(self, kernel_id):
self.kernel_id = kernel_id
self.session = Session()
self.save_on_message = self.on_message
self.on_message = self.on_first_message
def get_current_user(self):
password = self.get_secure_cookie("password")
if password is None:
# clear cookies, to prevent future Invalid cookie signature warnings
self._cookies = Cookie.SimpleCookie()
if self.application.password and self.application.password != password:
return None
return self.get_secure_cookie("user") or 'anonymous'
def _inject_cookie_message(self, msg):
"""Inject the first message, which is the document cookie,
for authentication."""
if isinstance(msg, unicode):
# Cookie can't constructor doesn't accept unicode strings for some reason
msg = msg.encode('utf8', 'replace')
try:
self._cookies = Cookie.SimpleCookie(msg)
except:
logging.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:
logging.warn("Couldn't authenticate WebSocket connection")
raise web.HTTPError(403)
self.on_message = self.save_on_message
class IOPubHandler(ZMQStreamHandler): class IOPubHandler(AuthenticatedZMQStreamHandler):
def initialize(self, *args, **kwargs): def initialize(self, *args, **kwargs):
self._kernel_alive = True self._kernel_alive = True
self._beating = False self._beating = False
self.iopub_stream = None self.iopub_stream = None
self.hb_stream = None self.hb_stream = None
def open(self, kernel_id): def on_first_message(self, msg):
try:
super(IOPubHandler, self).on_first_message(msg)
except web.HTTPError:
self.close()
return
km = self.application.kernel_manager km = self.application.kernel_manager
self.kernel_id = kernel_id
self.session = Session()
self.time_to_dead = km.time_to_dead self.time_to_dead = km.time_to_dead
kernel_id = self.kernel_id
try: try:
self.iopub_stream = km.create_iopub_stream(kernel_id) self.iopub_stream = km.create_iopub_stream(kernel_id)
self.hb_stream = km.create_hb_stream(kernel_id) self.hb_stream = km.create_hb_stream(kernel_id)
@ -187,9 +228,13 @@ class IOPubHandler(ZMQStreamHandler):
# close the connection. # close the connection.
if not self.stream.closed(): if not self.stream.closed():
self.stream.close() self.stream.close()
self.close()
else: else:
self.iopub_stream.on_recv(self._on_zmq_reply) self.iopub_stream.on_recv(self._on_zmq_reply)
self.start_hb(self.kernel_died) self.start_hb(self.kernel_died)
def on_message(self, msg):
pass
def on_close(self): def on_close(self):
# This method can be called twice, once by self.kernel_died and once # This method can be called twice, once by self.kernel_died and once
@ -245,15 +290,20 @@ class IOPubHandler(ZMQStreamHandler):
self.on_close() self.on_close()
class ShellHandler(ZMQStreamHandler): class ShellHandler(AuthenticatedZMQStreamHandler):
def initialize(self, *args, **kwargs): def initialize(self, *args, **kwargs):
self.shell_stream = None self.shell_stream = None
def open(self, kernel_id): def on_first_message(self, msg):
try:
super(ShellHandler, self).on_first_message(msg)
except web.HTTPError:
self.close()
return
km = self.application.kernel_manager km = self.application.kernel_manager
self.max_msg_size = km.max_msg_size self.max_msg_size = km.max_msg_size
self.kernel_id = kernel_id kernel_id = self.kernel_id
try: try:
self.shell_stream = km.create_shell_stream(kernel_id) self.shell_stream = km.create_shell_stream(kernel_id)
except web.HTTPError: except web.HTTPError:
@ -261,8 +311,8 @@ class ShellHandler(ZMQStreamHandler):
# close the connection. # close the connection.
if not self.stream.closed(): if not self.stream.closed():
self.stream.close() self.stream.close()
self.close()
else: else:
self.session = Session()
self.shell_stream.on_recv(self._on_zmq_reply) self.shell_stream.on_recv(self._on_zmq_reply)
def on_message(self, msg): def on_message(self, msg):

View File

@ -91,6 +91,12 @@ var IPython = (function (IPython) {
console.log("Starting WS:", ws_url); console.log("Starting WS:", ws_url);
this.shell_channel = new this.WebSocket(ws_url + "/shell"); this.shell_channel = new this.WebSocket(ws_url + "/shell");
this.iopub_channel = new this.WebSocket(ws_url + "/iopub"); this.iopub_channel = new this.WebSocket(ws_url + "/iopub");
send_cookie = function(){
this.send(document.cookie);
console.log(this);
}
this.shell_channel.onopen = send_cookie;
this.iopub_channel.onopen = send_cookie;
}; };