Merge "console: make websocketproxy handles token from path"

2 files changed, 99 insertions, 3 deletions

diff --git a/nova/console/websocketproxy.py b/nova/console/websocketproxy.py
index fbc2be8933..ef684f56a2 100644
--- a/nova/console/websocketproxy.py
+++ b/nova/console/websocketproxy.py
@@ -20,6 +20,7 @@ Leverages websockify.py by Joel Martin
 
 import Cookie
 import socket
+import urlparse
 
 import websockify
 
@@ -39,9 +40,21 @@ class NovaProxyRequestHandlerBase(object):
         from eventlet import hubs
         hubs.use_hub()
 
-        cookie = Cookie.SimpleCookie()
-        cookie.load(self.headers.getheader('cookie'))
-        token = cookie['token'].value
+        # The nova expected behavior is to have token
+        # passed to the method GET of the request
+        query = urlparse.urlparse(self.path).query
+        token = urlparse.parse_qs(query).get("token", [""]).pop()
+        if not token:
+            # NoVNC uses it's own convention that forward token
+            # from the request to a cookie header, we should check
+            # also for this behavior
+            hcookie = self.headers.getheader('cookie')
+            if hcookie:
+                cookie = Cookie.SimpleCookie()
+                cookie.load(hcookie)
+                if 'token' in cookie:
+                    token = cookie['token'].value
+
         ctxt = context.get_admin_context()
         rpcapi = consoleauth_rpcapi.ConsoleAuthAPI()
         connect_info = rpcapi.check_token(ctxt, token=token)
diff --git a/nova/tests/console/test_websocketproxy.py b/nova/tests/console/test_websocketproxy.py
new file mode 100644
index 0000000000..1e51a4d5ec
--- /dev/null
+++ b/nova/tests/console/test_websocketproxy.py
@@ -0,0 +1,83 @@
+# 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. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+"""Tests for nova websocketproxy."""
+
+
+import mock
+
+from nova.console import websocketproxy
+from nova import test
+
+
+class NovaProxyRequestHandlerBaseTestCase(test.TestCase):
+
+    def setUp(self):
+        super(NovaProxyRequestHandlerBaseTestCase, self).setUp()
+
+        self.wh = websocketproxy.NovaProxyRequestHandlerBase()
+        self.wh.socket = mock.MagicMock()
+        self.wh.msg = mock.MagicMock()
+        self.wh.do_proxy = mock.MagicMock()
+        self.wh.headers = mock.MagicMock()
+
+    @mock.patch('nova.consoleauth.rpcapi.ConsoleAuthAPI.check_token')
+    def test_new_websocket_client(self, check_token):
+        check_token.return_value = {
+            'host': 'node1',
+            'port': '10000'
+        }
+        self.wh.socket.return_value = '<socket>'
+        self.wh.path = "ws://127.0.0.1/?token=123-456-789"
+
+        self.wh.new_websocket_client()
+
+        check_token.assert_called_with(mock.ANY, token="123-456-789")
+        self.wh.socket.assert_called_with('node1', 10000, connect=True)
+        self.wh.do_proxy.assert_called_with('<socket>')
+
+    @mock.patch('nova.consoleauth.rpcapi.ConsoleAuthAPI.check_token')
+    def test_new_websocket_client_token_invalid(self, check_token):
+        check_token.return_value = False
+
+        self.wh.path = "ws://127.0.0.1/?token=XXX"
+
+        self.assertRaises(Exception, self.wh.new_websocket_client)  # noqa
+        check_token.assert_called_with(mock.ANY, token="XXX")
+
+    @mock.patch('nova.consoleauth.rpcapi.ConsoleAuthAPI.check_token')
+    def test_new_websocket_client_novnc(self, check_token):
+        check_token.return_value = {
+            'host': 'node1',
+            'port': '10000'
+        }
+        self.wh.socket.return_value = '<socket>'
+        self.wh.path = "http://127.0.0.1/"
+        self.wh.headers.getheader.return_value = "token=123-456-789"
+
+        self.wh.new_websocket_client()
+
+        check_token.assert_called_with(mock.ANY, token="123-456-789")
+        self.wh.socket.assert_called_with('node1', 10000, connect=True)
+        self.wh.do_proxy.assert_called_with('<socket>')
+
+    @mock.patch('nova.consoleauth.rpcapi.ConsoleAuthAPI.check_token')
+    def test_new_websocket_client_novnc_token_invalid(self, check_token):
+        check_token.return_value = False
+
+        self.wh.path = "http://127.0.0.1/"
+        self.wh.headers.getheader.return_value = "token=XXX"
+
+        self.assertRaises(Exception, self.wh.new_websocket_client)  # noqa
+        check_token.assert_called_with(mock.ANY, token="XXX")