feat(websocket): add authentication handshake to websocket connection

If we are already logged in, or have just successfully logged in, we need to authenticate our connection. This will let the backend know that it should send the events associated with the current session's profileID over the websocket connection.

If we ever logout, we need to unsubscribe to this connection.

src/features/app/App.tsx

@@ -1,6 +1,6 @@
 import React, { useEffect } from 'react'
 import { ConnectedRouter } from 'connected-react-router'
-import { useDispatch } from 'react-redux'
+import { useDispatch, useSelector } from 'react-redux'
 
 import { history } from '../../store/store'
 import Routes from '../../routes'
@@ -10,13 +10,15 @@ import SnackBar from '../snackBar/SnackBar'
 import ToastRenderer from '../toast/ToastRenderer'
 
 import './App.css'
+import { selectSessionTokens } from '../session/state/sessionState'
 
 const App: React.FC<any> = () => {
   const dispatch = useDispatch()
 
+  const tokens = useSelector(selectSessionTokens)
   useEffect(() => {
-    dispatch(wsConnect())
-  })
+    dispatch(wsConnect(tokens.token))
+  }, [dispatch, tokens.token])
 
   return (
     <div id='app' className='flex flex-col h-screen w-screen' style={{ minWidth: '1100px' }}>

src/features/websocket/middleware/websocket.ts

@@ -33,6 +33,10 @@ import {
   ETAutomationDeploySaveDatasetEnd,
   ETAutomationDeploySaveWorkflowStart,
   ETAutomationDeploySaveWorkflowEnd,
+  WSSubscribeSuccess,
+  WSSubscribeFailure,
+  WSSubscribeRequest,
+  WSUnsubscribeRequest,
 } from '../../../qri/events'
 import { wsConnectionChange } from '../state/websocketActions'
 import { WS_CONNECT, WS_DISCONNECT } from '../state/websocketState'
@@ -56,8 +60,8 @@ export type RemoteEvents = Record<string, RemoteEvent>
 
 const WEBSOCKETS_URL = process.env.REACT_APP_WS_BASE_URL || 'ws://localhost:2503'
 const WEBSOCKETS_PROTOCOL = 'qri-websocket'
-const numReconnectAttempts = 2
-const msToAddBeforeReconnectAttempt = 3000
+const numReconnectAttempts: number = 2
+const msToAddBeforeReconnectAttempt: number = 3000
 
 function newReconnectDeadline(): Date {
   let d = new Date()
@@ -71,17 +75,18 @@ const middleware = () => {
     status: WSConnectionStatus.disconnected,
   }
 
-  const onOpen = (dispatch: Dispatch ) => (event: Event) => {
+  const onOpen = (dispatch: Dispatch, token: string) => (event: Event) => {
     state = {
       status: WSConnectionStatus.connected,
       reconnectAttemptsRemaining: 0,
       reconnectTime: undefined,
     }
     const stateCopy = NewWebsocketState(state.status)
     dispatch(wsConnectionChange(stateCopy))
+    subscribe(token)
   }
 
-  const onClose = (dispatch: Dispatch) => (event: Event) => {
+  const onClose = (dispatch: Dispatch, token: string) => (event: Event) => {
     // code 1006 is an "Abnormal Closure" where no close frame is sent. Happens
     // when there isn't a websocket host on the other end (aka: server down)
     // https://developer.mozilla.org/en-US/docs/Web/API/CloseEvent
@@ -96,7 +101,7 @@ const middleware = () => {
       state.reconnectAttemptsRemaining = numReconnectAttempts
     }
 
-    reconnect(dispatch)
+    reconnect(dispatch, token)
     const stateCopy = NewWebsocketState(state.status, state.reconnectAttemptsRemaining, state.reconnectTime)
     dispatch(wsConnectionChange(stateCopy))
   }
@@ -173,6 +178,12 @@ const middleware = () => {
         case ETAutomationDeploySaveDatasetEnd:
           dispatch(deploySaveDatasetEnded(event.data, event.sessionID))
           break
+        case WSSubscribeSuccess:
+          console.log("successfully authenticated websocket connection")
+          break
+        case WSSubscribeFailure:
+          console.log("unable to authenticate websocket connection", event.error)
+          break
 
         default:
           // console.log(`received websocket event: ${event.type}`)
@@ -183,26 +194,66 @@ const middleware = () => {
     }
   }
 
-  const connect = (dispatch: Dispatch) => {
+  const connect = (dispatch: Dispatch, token: string) => {
     if (socket !== undefined) {
       socket.close()
     }
     // connect to the remote host
     socket = new WebSocket(WEBSOCKETS_URL, WEBSOCKETS_PROTOCOL)
     socket.onmessage = onMessage(dispatch)
-    socket.onclose = onClose(dispatch)
-    socket.onopen = onOpen(dispatch)
+    socket.onclose = onClose(dispatch, token)
+    socket.onopen = onOpen(dispatch, token)
+  }
+
+  const subscribe = (token: string) => {
+    if (token === "") {
+      return 
+    }
+    if (socket === undefined) {
+      return
+    }
+    var reconnectAttemptsRemaining = numReconnectAttempts
+    const msg = JSON.stringify({ type: WSSubscribeRequest, payload: { token }})
+    const attempt = setInterval(() => {
+      if (reconnectAttemptsRemaining === 0) {
+        clearInterval(attempt)
+      }
+      if (socket?.readyState === 1) {
+        socket.send(msg)
+        clearInterval(attempt)
+        return
+      }
+      reconnectAttemptsRemaining--
+      }, msToAddBeforeReconnectAttempt)
+  }
+
+  const unsubscribe = () => {
+    if (socket === undefined) {
+      return
+    }
+    var reconnectAttemptsRemaining = numReconnectAttempts
+    const attempt = setInterval(() => {
+      if (reconnectAttemptsRemaining === 0) {
+        clearInterval(attempt)
+      }
+      if (socket?.readyState === 1) {
+        socket.send(JSON.stringify({ type: WSUnsubscribeRequest }))
+        clearInterval(attempt)
+        return
+      }
+      reconnectAttemptsRemaining--
+      }, msToAddBeforeReconnectAttempt)
   }
 
-  const reconnect = (dispatch: Dispatch) => {
+  const reconnect = (dispatch: Dispatch, token: string) => {
     if (state.reconnectAttemptsRemaining && state.reconnectAttemptsRemaining > 0) {
       state = {
         status: WSConnectionStatus.interrupted,
         reconnectAttemptsRemaining: state.reconnectAttemptsRemaining - 1,
         reconnectTime: newReconnectDeadline()
       }
       setTimeout(() => {
-        connect(dispatch)
+        connect(dispatch, token)
       }, msToAddBeforeReconnectAttempt)
       return
     }
@@ -217,14 +268,19 @@ const middleware = () => {
   return (store: Store<RootState>) => (next: Dispatch<AnyAction>) => (action: AnyAction) => {
     switch (action.type) {
       case WS_CONNECT:
-        connect(next)
+        connect(next, action.token)
         break
       case WS_DISCONNECT:
         if (socket !== undefined) {
           socket.close()
         }
         socket = undefined
         break
+      case 'LOGIN_SUCCESS':
+        subscribe(action.token)
+        break
+      case 'LOGOUT_SUCCESS':
+        unsubscribe()
     }
 
     return next(action)

src/features/websocket/state/websocketActions.ts

@@ -6,9 +6,13 @@ import {
   WebsocketState,
 } from './websocketState'
 
+interface ConnectAction extends Action {
+  token: string
+}
+
 // tell websocket to connect
-export function wsConnect(): Action {
-  return { type: WS_CONNECT }
+export function wsConnect(token?: string): ConnectAction {
+  return { type: WS_CONNECT, token: token || '' }
 }
 
 // tell websocket to disconnect

src/qri/events.ts

@@ -51,3 +51,10 @@ export const ETAutomationDeploySaveWorkflowEnd = "automation:DeploySaveWorkflowE
 
 export const ETAutomationDeploySaveDatasetStart = "automation:DeploySaveDatasetStart"
 export const ETAutomationDeploySaveDatasetEnd = "automation:DeploySaveDatasetEnd"
+
+// Websocket message types describe the different types of messages that can be
+// sent over the websocket connection to establish the authentication handshake
+export const WSSubscribeRequest = "subscribe:request"
+export const WSSubscribeSuccess = "subscribe:success"
+export const WSSubscribeFailure = "subscribe:failure"
+export const WSUnsubscribeRequest = "unsubscribe:request"