From 505fb260610e557e27bbc5d27515337ea07e0e3e Mon Sep 17 00:00:00 2001
From: Henry Jameson <me@hjkos.com>
Date: Sun, 8 Dec 2019 19:18:38 +0200
Subject: [PATCH] better wrapper for websocket

---
 src/modules/api.js                            | 43 +++++++++--------
 src/services/api/api.service.js               | 46 ++++++++++++++++++-
 .../backend_interactor_service.js             | 15 ++----
 3 files changed, 69 insertions(+), 35 deletions(-)

diff --git a/src/modules/api.js b/src/modules/api.js
index 185c9db6..593f8498 100644
--- a/src/modules/api.js
+++ b/src/modules/api.js
@@ -34,36 +34,35 @@ const api = {
     // MastoAPI 'User' sockets
     startMastoUserSocket (store) {
       const { state, dispatch } = store
-      state.mastoUserSocket = state.backendInteractor
-        .startUserSocket({
-          store,
-          onMessage: (message) => {
-            if (!message) return // pings
-            if (message.event === 'notification') {
-              dispatch('addNewNotifications', {
-                notifications: [message.notification],
-                older: false
-              })
-            } else if (message.event === 'update') {
-              dispatch('addNewStatuses', {
-                statuses: [message.status],
-                userId: false,
-                showImmediately: false,
-                timeline: 'friends'
-              })
-            }
+      state.mastoUserSocket = state.backendInteractor.startUserSocket({ store })
+      state.mastoUserSocket.addEventListener(
+        'message',
+        ({ detail: message }) => {
+          if (!message) return // pings
+          if (message.event === 'notification') {
+            dispatch('addNewNotifications', {
+              notifications: [message.notification],
+              older: false
+            })
+          } else if (message.event === 'update') {
+            dispatch('addNewStatuses', {
+              statuses: [message.status],
+              userId: false,
+              showImmediately: false,
+              timeline: 'friends'
+            })
           }
-        })
-      state.mastoUserSocket.addEventListener('error', error => {
+        }
+      )
+      state.mastoUserSocket.addEventListener('error', ({ detail: error }) => {
         console.error('Error in MastoAPI websocket:', error)
       })
-      state.mastoUserSocket.addEventListener('close', closeEvent => {
+      state.mastoUserSocket.addEventListener('close', ({ detail: closeEvent }) => {
         const ignoreCodes = new Set([
           1000, // Normal (intended) closure
           1001 // Going away
         ])
         const { code } = closeEvent
-        console.debug('Socket closure event:', closeEvent)
         if (ignoreCodes.has(code)) {
           console.debug(`Not restarting socket becasue of closure code ${code} is in ignore list`)
         } else {
diff --git a/src/services/api/api.service.js b/src/services/api/api.service.js
index 7f27564f..c6818df4 100644
--- a/src/services/api/api.service.js
+++ b/src/services/api/api.service.js
@@ -953,8 +953,52 @@ const MASTODON_STREAMING_EVENTS = new Set([
   'filters_changed'
 ])
 
+// A thin wrapper around WebSocket API that allows adding a pre-processor to it
+// Uses EventTarget and a CustomEvent to proxy events
+export const ProcessedWS = ({
+  url,
+  preprocessor = handleMastoWS,
+  id = 'Unknown'
+}) => {
+  const eventTarget = new EventTarget()
+  const socket = new WebSocket(url)
+  if (!socket) throw new Error(`Failed to create socket ${id}`)
+  const proxy = (original, eventName, processor = a => a) => {
+    original.addEventListener(eventName, (eventData) => {
+      eventTarget.dispatchEvent(new CustomEvent(
+        eventName,
+        { detail: processor(eventData) }
+      ))
+    })
+  }
+  socket.addEventListener('open', (wsEvent) => {
+    console.debug(`[WS][${id}] Socket connected`, wsEvent)
+  })
+  socket.addEventListener('error', (wsEvent) => {
+    console.debug(`[WS][${id}] Socket errored`, wsEvent)
+  })
+  socket.addEventListener('close', (wsEvent) => {
+    console.debug(
+      `[WS][${id}] Socket disconnected with code ${wsEvent.code}`,
+      wsEvent
+    )
+  })
+  socket.addEventListener('message', (wsEvent) => {
+    console.debug(
+      `[WS][${id}] Message received`,
+      wsEvent
+    )
+  })
+
+  proxy(socket, 'open')
+  proxy(socket, 'close')
+  proxy(socket, 'message', preprocessor)
+  proxy(socket, 'error')
+
+  return eventTarget
+}
+
 export const handleMastoWS = (wsEvent) => {
-  console.debug('Event', wsEvent)
   const { data } = wsEvent
   if (!data) return
   const parsedEvent = JSON.parse(data)
diff --git a/src/services/backend_interactor_service/backend_interactor_service.js b/src/services/backend_interactor_service/backend_interactor_service.js
index 33b79a40..b7372ed0 100644
--- a/src/services/backend_interactor_service/backend_interactor_service.js
+++ b/src/services/backend_interactor_service/backend_interactor_service.js
@@ -1,4 +1,4 @@
-import apiService, { getMastodonSocketURI, handleMastoWS } from '../api/api.service.js'
+import apiService, { getMastodonSocketURI, ProcessedWS } from '../api/api.service.js'
 import timelineFetcherService from '../timeline_fetcher/timeline_fetcher.service.js'
 import notificationsFetcher from '../notifications_fetcher/notifications_fetcher.service.js'
 import followRequestFetcher from '../../services/follow_request_fetcher/follow_request_fetcher.service'
@@ -20,19 +20,10 @@ const backendInteractorService = credentials => ({
     return followRequestFetcher.startFetching({ store, credentials })
   },
 
-  startUserSocket ({ store, onMessage }) {
+  startUserSocket ({ store }) {
     const serv = store.rootState.instance.server.replace('http', 'ws')
     const url = serv + getMastodonSocketURI({ credentials, stream: 'user' })
-    const socket = new WebSocket(url)
-    console.debug('Socket created:', socket)
-    if (socket) {
-      socket.addEventListener('open', (wsEvent) => console.debug('MastoAPI User WebSocket connection established'))
-      socket.addEventListener('message', (wsEvent) => onMessage(handleMastoWS(wsEvent)))
-      socket.addEventListener('error', (error) => console.error('MastoApi User WebSocket Error:', error))
-      return socket
-    } else {
-      throw new Error('failed to connect to socket')
-    }
+    return ProcessedWS({ url, id: 'User' })
   },
 
   ...Object.entries(apiService).reduce((acc, [key, func]) => {