Merge pull request #1851 from hypothesis/ws-origin-check

Refactor WebSocket origin security

h/static/scripts/app.coffee

@@ -57,12 +57,6 @@ configure = [
     # TODO: move all front-end templates into their own directory for safety
     basePattern = baseURI.replace /\/[^\/]*$/, '/**.html'
     $sceDelegateProvider.resourceUrlWhitelist ['self', basePattern]
-
-    streamerProvider.urlFn = ['xsrf', (xsrf) ->
-      base = baseURI.replace(/^http/, 'ws')
-      "#{base}ws?csrf_token=#{xsrf.token}"
-    ]
 ]
 
-
 angular.module('h', imports, configure)

h/static/scripts/controllers.coffee

 class AppController
   this.$inject = [
-    '$location', '$route', '$scope', '$timeout',
+    '$location', '$route', '$scope', '$timeout', '$window',
     'annotator', 'auth', 'documentHelpers', 'drafts', 'flash', 'identity',
     'permissions', 'streamer', 'streamfilter'
-
   ]
   constructor: (
-     $location,   $route,   $scope,   $timeout,
+     $location,   $route,   $scope,   $timeout,   $window,
      annotator,   auth,   documentHelpers,   drafts,   flash,   identity,
      permissions,   streamer,   streamfilter,
 
@@ -15,6 +14,7 @@ class AppController
 
     $scope.auth = auth
     isFirstRun = $location.search().hasOwnProperty('firstrun')
+    streamerUrl = documentHelpers.baseURI.replace(/^http/, 'ws') + 'ws'
 
     applyUpdates = (action, data) ->
       # Update the application with new data from the websocket.
@@ -110,8 +110,9 @@ class AppController
 
       # Reload services
       initStore()
+
       streamer.close()
-      streamer.open()
+      streamer.open($window.WebSocket, streamerUrl)
 
     $scope.$watch 'socialView.name', (newValue, oldValue) ->
       return if newValue is oldValue

h/static/scripts/streamer-service.coffee

-ST_CLOSED = 1
-ST_CONNECTING = 2
-ST_OPEN = 3
+ST_CONNECTING = 0
+ST_OPEN = 1
+ST_CLOSING = 2
+ST_CLOSED = 3
 
 ###*
 # @ngdoc service
 # @name Streamer
 #
-# @param {String} urlFn A function that will be called with injections to
-# generate the socket URL.
+# @property {string} clientId A unique identifier for this client.
 #
 # @description
 # Provides access to the streamer websocket.
 ###
 class Streamer
-  constructor: (transport, urlFn) ->
-    this.onmessage = ->
+  constructor: ->
+    this.clientId = null
+
+    this.onopen = angular.noop
+    this.onclose = angular.noop
+    this.onmessage = angular.noop
 
     this._failCount = 0
     this._queue = []
     this._state = ST_CLOSED
-    this._transport = transport
-    this._urlFn = urlFn
 
   ###*
   # @ngdoc method
   # @name Streamer#open
   #
+  # @param {Object} transport The transport class to create.
+  # @param {string} url The URL to which to connect.
+  # @param {Object|Array} protocols The protocol (or protocols) to use.
+  #
   # @description
   # Open the streamer websocket if it is not already open or connecting. Handles
   # connection failures and sets up onmessage handlers.
   ###
-  open: ->
+  open: (transport, url, protocols) ->
     if this._state == ST_OPEN || this._state == ST_CONNECTING
       return
 
     self = this
-    this._sock = new this._transport(this._urlFn())
+
+    if protocols
+      this._sock = new transport(url, protocols)
+    else
+      this._sock = new transport(url)
+
     this._state = ST_CONNECTING
 
     this._sock.onopen = ->
       self._state = ST_OPEN
       self._failCount = 0
 
-      clientId = uuid.v4()
-      setAjaxClientId(clientId)
+      # Send the client id
+      if self.clientId
+        self.send(messageType: 'client_id', value: self.clientId)
+
+      # Give the application a chance to initialize the connection
+      self.onopen(name: 'open')
 
-      # Generate and send our client ID
-      self.send({
-        messageType: 'client_id'
-        value: clientId
-      })
       # Process queued messages
       self._sendQueue()
 
@@ -57,7 +67,9 @@ class Streamer
       self._state = ST_CLOSED
       self._failCount++
 
-      setTimeout((-> self.open()), backoff(self._failCount, 10))
+      reconnect = angular.bind(self, self.open, transport, url, protocols)
+      waitFor = backoff(self._failCount, 10)
+      setTimeout(reconnect, waitFor)
 
     this._sock.onmessage = (msg) ->
       self.onmessage(JSON.parse(msg.data))
@@ -70,13 +82,18 @@ class Streamer
   # Close the streamer socket.
   ###
   close: ->
-    if this._state == ST_CLOSED
+    if this._state == ST_CLOSING or this._state == ST_CLOSED
       return
 
+    self = this
+
     this._sock.onclose = ->
+      self._state = ST_CLOSED
+      self.onclose()
+
     this._sock.close()
     this._sock = null
-    this._state = ST_CLOSED
+    this._state = ST_CLOSING
 
   ###*
   # @ngdoc method
@@ -104,22 +121,16 @@ backoff = (index, max) ->
   index = Math.min(index, max)
   return 500 * Math.random() * (Math.pow(2, index) - 1)
 
-setAjaxClientId = (clientId) ->
-  $.ajaxSetup({
-    headers: {
-      'X-Client-Id': clientId
-    }
-  })
-
-streamerProvider = ->
-  provider = {}
-  provider.urlFn = null
-  provider.$get = ['$injector', '$window', ($injector, $window) ->
-    urlFn = angular.bind $injector, $injector.invoke, provider.urlFn
-    new Streamer($window.WebSocket, urlFn)
-  ]
-  return provider
 
+run = [
+  '$http', '$window', 'streamer'
+  ($http,   $window,   streamer) ->
+    clientId = uuid.v4()
+    streamer.clientId = clientId
+    $.ajaxSetup(headers: {'X-Client-Id': clientId})
+    $http.defaults.headers.common['X-Client-Id'] = clientId
+]
 
 angular.module('h.streamer', [])
-.provider('streamer', streamerProvider)
+.service('streamer', Streamer)
+.run(run)

tests/js/streamer-service-test.coffee

@@ -6,19 +6,16 @@ describe 'streamer', ->
   WebSocket = null
   fakeSock = null
   streamer = null
+  url = 'wss://magicstreemz/giraffe'
 
   beforeEach module('h.streamer')
 
-  beforeEach module ($provide, streamerProvider) ->
+  beforeEach module ->
     fakeSock = {
       send: sandbox.spy()
       close: sandbox.spy()
     }
     WebSocket = sandbox.stub().returns(fakeSock)
-    $provide.decorator '$window', ($delegate) ->
-      angular.extend $delegate, {WebSocket}
-    $provide.value 'webSocketUrl', 'wss://magicstreemz/giraffe'
-    streamerProvider.urlFn = (webSocketUrl) -> webSocketUrl
     return
 
   beforeEach inject (_streamer_) ->
@@ -28,22 +25,35 @@ describe 'streamer', ->
     sandbox.restore()
 
   it 'calls the transport function with the new keyword', ->
-    streamer.open()
+    streamer.open(WebSocket, url)
+
     assert.calledWithNew(WebSocket)
 
   it 'creates a socket with the correct base URL', ->
-    streamer.open()
+    streamer.open(WebSocket, url)
 
     assert.calledWith(WebSocket, 'wss://magicstreemz/giraffe')
 
-  it 'does not open another socket while a socket is connecting', ->
-    streamer.open()
-    streamer.open()
+  it 'does not open another socket while connecting or connected', ->
+    streamer.open(WebSocket, url)
+    streamer.open(WebSocket, url)
+
+    assert.calledOnce(WebSocket)
+
+    fakeSock.onopen()
+    streamer.open(WebSocket, url)
 
     assert.calledOnce(WebSocket)
 
-  it 'queues messages until the socket is open', ->
-    streamer.open()
+  it 'does not close the socket again when already closing', ->
+    streamer.open(WebSocket, url)
+    streamer.close()
+    streamer.close()
+
+    assert.calledOnce(fakeSock.close)
+
+  it 'queues messages kuntil the socket is open', ->
+    streamer.open(WebSocket, url)
     streamer.send({animal: 'elephant'})
 
     assert.notCalled(fakeSock.send)
@@ -52,8 +62,18 @@ describe 'streamer', ->
 
     assert.called(fakeSock.send)
 
+  it 'calls the onopen handler once the socket is open', ->
+    streamer.onopen = sinon.spy()
+    streamer.open(WebSocket, url)
+
+    assert.notCalled(streamer.onopen)
+
+    fakeSock.onopen()
+
+    assert.called(streamer.onopen)
+
   it 'preserves message ordering in the queue', ->
-    streamer.open()
+    streamer.open(WebSocket, url)
     streamer.send({animal: 'elephant'})
     streamer.send({animal: 'giraffe'})
     fakeSock.onopen()
@@ -64,7 +84,7 @@ describe 'streamer', ->
     assert.equal(secondAnimal, 'giraffe')
 
   it 'converts message data to JSON', ->
-    streamer.open()
+    streamer.open(WebSocket, url)
     streamer.send({animal: 'badger'})
     fakeSock.onopen()
 
@@ -72,7 +92,7 @@ describe 'streamer', ->
 
   it 'sends a client ID as the first message once the socket opens', ->
     streamer.send({animal: 'elephant'})
-    streamer.open()
+    streamer.open(WebSocket, url)
     fakeSock.onopen()
 
     msg = fakeSock.send.getCall(0).args[0]
@@ -83,21 +103,22 @@ describe 'streamer', ->
   it 'attempts to reopen the socket on connection failure', ->
     clock = sandbox.useFakeTimers()
 
-    streamer.open()
+    streamer.open(WebSocket, url)
     fakeSock.onclose()
 
     clock.tick(500)
 
     assert.calledTwice(WebSocket)
+    assert.match(WebSocket.getCall(0).args, WebSocket.getCall(1).args)
 
   it 'closes the socket when close is called', ->
-    streamer.open()
+    streamer.open(WebSocket, url)
     streamer.close()
 
     assert.calledOnce(fakeSock.close)
 
   it 'only closes the socket once', ->
-    streamer.open()
+    streamer.open(WebSocket, url)
     streamer.close()
     streamer.close()
 
@@ -105,7 +126,7 @@ describe 'streamer', ->
 
   it 'does not try and reopen the socket when closed explicitly', ->
     clock = sandbox.useFakeTimers()
-    streamer.open()
+    streamer.open(WebSocket, url)
     streamer.close()
     fakeSock.onclose()
 
@@ -114,12 +135,12 @@ describe 'streamer', ->
 
   it 'calls the onmessage handler when the socket receives a message', ->
     streamer.onmessage = sinon.spy()
-    streamer.open()
+    streamer.open(WebSocket, url)
     fakeSock.onmessage(data: JSON.stringify({animal: 'baboon'}))
     assert.called(streamer.onmessage)
 
   it 'calls the onmessage handler with parsed JSON', ->
     streamer.onmessage = sinon.spy()
-    streamer.open()
+    streamer.open(WebSocket, url)
     fakeSock.onmessage(data: JSON.stringify({animal: 'baboon'}))
     assert.calledWith(streamer.onmessage, {animal: 'baboon'})