still partial

Signed-off-by: Andrew Thornton <[email protected]>

Author: zeripath

modules/context/response.go

@@ -84,14 +84,29 @@ func (r *Response) Before(f func(ResponseWriter)) {
 	r.befores = append(r.befores, f)
 }
 
+type hijackerResponse struct {
+	*Response
+	http.Hijacker
+}
+
 // NewResponse creates a response
-func NewResponse(resp http.ResponseWriter) *Response {
+func NewResponse(resp http.ResponseWriter) ResponseWriter {
 	if v, ok := resp.(*Response); ok {
 		return v
 	}
-	return &Response{
+	hijacker, ok := resp.(http.Hijacker)
+
+	response := &Response{
 		ResponseWriter: resp,
 		status:         0,
 		befores:        make([]func(ResponseWriter), 0),
 	}
+	if ok {
+		return hijackerResponse{
+			Response: response,
+			Hijacker: hijacker,
+		}
+	}
+
+	return response
 }

routers/web/events/websocket.go

@@ -5,11 +5,15 @@
 package events
 
 import (
+	"net/http"
+	"net/url"
+
 	"code.gitea.io/gitea/modules/context"
 	"code.gitea.io/gitea/modules/eventsource"
 	"code.gitea.io/gitea/modules/graceful"
 	"code.gitea.io/gitea/modules/json"
 	"code.gitea.io/gitea/modules/log"
+	"code.gitea.io/gitea/modules/setting"
 	"code.gitea.io/gitea/routers/web/auth"
 	"github.com/gorilla/websocket"
 )
@@ -25,8 +29,28 @@ func Websocket(ctx *context.Context) {
 	upgrader := websocket.Upgrader{
 		ReadBufferSize:  1024,
 		WriteBufferSize: 1024,
+		CheckOrigin: func(r *http.Request) bool {
+			origin := r.Header["Origin"]
+			if len(origin) == 0 {
+				return true
+			}
+			u, err := url.Parse(origin[0])
+			if err != nil {
+				return false
+			}
+			appURLURL, err := url.Parse(setting.AppURL)
+			if err != nil {
+				return true
+			}
+
+			return u.Host == appURLURL.Host
+		},
 	}
 
+	// Because http proxies will tend not to pass these headers
+	ctx.Req.Header.Add("Upgrade", "websocket")
+	ctx.Req.Header.Add("Connection", "Upgrade")
+
 	conn, err := upgrader.Upgrade(ctx.Resp, ctx.Req, nil)
 	if err != nil {
 		log.Error("Unable to upgrade due to error: %v", err)

web_src/js/features/eventsource.sharedworker.js

@@ -46,9 +46,13 @@ class Source {
     if (this.listening[eventType]) return;
     this.listening[eventType] = true;
     this.eventSource.addEventListener(eventType, (event) => {
+      let data;
+      if (event.data) {
+        data = JSON.parse(event.data);
+      }
       this.notifyClients({
         type: eventType,
-        data: event.data
+        data
       });
     });
   }

web_src/js/features/notification.js

@@ -24,10 +24,9 @@ export function initNotificationsTable() {
   });
 }
 
-async function receiveUpdateCount(event) {
+async function receiveUpdateCount(data) {
+  console.log(data);
   try {
-    const data = JSON.parse(event.data);
-
     const notificationCount = document.querySelector('.notification_count');
     if (data.Count > 0) {
       notificationCount.classList.remove('hidden');
@@ -36,9 +35,10 @@ async function receiveUpdateCount(event) {
     }
 
     notificationCount.textContent = `${data.Count}`;
+    console.log(notificationCount);
     await updateNotificationTable();
   } catch (error) {
-    console.error(error, event);
+    console.error(error, data);
   }
 }
 
@@ -49,9 +49,20 @@ export function initNotificationCount() {
     return;
   }
 
-  if (notificationSettings.EventSourceUpdateTime > 0 && !!window.EventSource && window.SharedWorker) {
+  let worker;
+  let workerUrl;
+
+  if (notificationSettings.EventSourceUpdateTime > 0 && !!window.WebSocket && window.SharedWorker) {
+    // Try to connect to the event source via the shared worker first
+    worker = new SharedWorker(`${__webpack_public_path__}js/websocket.sharedworker.js`, 'notification-worker');
+    workerUrl = `${window.location.origin}${appSubUrl}/user/websocket`;
+  } else if (notificationSettings.EventSourceUpdateTime > 0 && !!window.EventSource && window.SharedWorker) {
     // Try to connect to the event source via the shared worker first
-    const worker = new SharedWorker(`${__webpack_public_path__}js/eventsource.sharedworker.js`, 'notification-worker');
+    worker = new SharedWorker(`${__webpack_public_path__}js/eventsource.sharedworker.js`, 'notification-worker');
+    workerUrl = `${window.location.origin}${appSubUrl}/user/events`;
+  }
+
+  if (worker) {
     worker.addEventListener('error', (event) => {
       console.error(event);
     });
@@ -60,15 +71,16 @@ export function initNotificationCount() {
     });
     worker.port.postMessage({
       type: 'start',
-      url: `${window.location.origin}${appSubUrl}/user/events`,
+      url: workerUrl,
     });
     worker.port.addEventListener('message', (event) => {
       if (!event.data || !event.data.type) {
         console.error(event);
         return;
       }
+      console.log(event);
       if (event.data.type === 'notification-count') {
-        const _promise = receiveUpdateCount(event.data);
+        const _promise = receiveUpdateCount(event.data.data).then(console.log('done'));
       } else if (event.data.type === 'error') {
         console.error(event.data);
       } else if (event.data.type === 'logout') {

web_src/js/features/stopwatch.js

@@ -26,9 +26,19 @@ export function initStopwatch() {
     $(this).parent().trigger('submit');
   });
 
-  if (notificationSettings.EventSourceUpdateTime > 0 && !!window.EventSource && window.SharedWorker) {
+  let worker;
+  let workerUrl;
+  if (notificationSettings.EventSourceUpdateTime > 0 && !!window.WebSocket && window.SharedWorker) {
     // Try to connect to the event source via the shared worker first
-    const worker = new SharedWorker(`${__webpack_public_path__}js/eventsource.sharedworker.js`, 'notification-worker');
+    worker = new SharedWorker(`${__webpack_public_path__}js/websocket.sharedworker.js`, 'notification-worker');
+    workerUrl = `${window.location.origin}${appSubUrl}/user/websocket`;
+  } else if (notificationSettings.EventSourceUpdateTime > 0 && !!window.EventSource && window.SharedWorker) {
+    // Try to connect to the event source via the shared worker first
+    worker = new SharedWorker(`${__webpack_public_path__}js/eventsource.sharedworker.js`, 'notification-worker');
+    workerUrl = `${window.location.origin}${appSubUrl}/user/events`;
+  }
+
+  if (worker) {
     worker.addEventListener('error', (event) => {
       console.error(event);
     });
@@ -37,15 +47,15 @@ export function initStopwatch() {
     });
     worker.port.postMessage({
       type: 'start',
-      url: `${window.location.origin}${appSubUrl}/user/events`,
+      url: workerUrl,
     });
     worker.port.addEventListener('message', (event) => {
       if (!event.data || !event.data.type) {
         console.error(event);
         return;
       }
       if (event.data.type === 'stopwatches') {
-        updateStopwatchData(JSON.parse(event.data.data));
+        updateStopwatchData(event.data.data);
       } else if (event.data.type === 'error') {
         console.error(event.data);
       } else if (event.data.type === 'logout') {

web_src/js/features/websocket.sharedworker.js

@@ -0,0 +1,163 @@
+const sourcesByUrl = {};
+const sourcesByPort = {};
+
+class Source {
+  constructor(url) {
+    this.url = url.replace(/^http/, 'ws');
+    this.webSocket = new WebSocket(this.url);
+    this.listening = {};
+    this.clients = [];
+    this.listen('open');
+    this.listen('close');
+    this.listen('logout');
+    this.listen('notification-count');
+    this.listen('stopwatches');
+    this.listen('error');
+    this.webSocket.addEventListener('error', (error) => {
+      this.lastError = error;
+    });
+    this.webSocket.addEventListener('message', (event) => {
+      const message = JSON.parse(event.data);
+      if (!message) {
+        return;
+      }
+      if (this.listening[message.Name]) {
+        this.notifyClients({
+          type: message.Name,
+          data: message.Data
+        });
+      }
+    });
+  }
+
+  register(port) {
+    if (this.clients.includes(port)) return;
+
+    this.clients.push(port);
+
+    port.postMessage({
+      type: 'status',
+      message: `registered to ${this.url}`,
+    });
+
+    if (!this.webSocket) {
+      if (this.lastError) {
+        port.postMessage({
+          type: 'error',
+          message: `websocket disconnected: ${this.lastError}`
+        });
+      } else {
+        port.postMessage({
+          type: 'error',
+          message: 'websocket disconnected'
+        });
+      }
+    }
+  }
+
+  deregister(port) {
+    const portIdx = this.clients.indexOf(port);
+    if (portIdx < 0) {
+      return this.clients.length;
+    }
+    this.clients.splice(portIdx, 1);
+    return this.clients.length;
+  }
+
+  close() {
+    if (!this.webSocket) return;
+
+    this.webSocket.close();
+    this.webSocket = null;
+  }
+
+  listen(eventType) {
+    if (this.listening[eventType]) return;
+    this.listening[eventType] = true;
+    this.webSocket.addEventListener(eventType, (event) => {
+      this.notifyClients({
+        type: eventType,
+        data: event.data
+      });
+    });
+  }
+
+  notifyClients(event) {
+    for (const client of this.clients) {
+      client.postMessage(event);
+    }
+  }
+
+  status(port) {
+    port.postMessage({
+      type: 'status',
+      message: `url: ${this.url} readyState: ${this.webSocket.readyState}`,
+    });
+  }
+}
+
+self.addEventListener('connect', (e) => {
+  for (const port of e.ports) {
+    port.addEventListener('message', (event) => {
+      if (event.data.type === 'start') {
+        const url = event.data.url;
+        if (sourcesByUrl[url]) {
+          // we have a Source registered to this url
+          const source = sourcesByUrl[url];
+          source.register(port);
+          sourcesByPort[port] = source;
+          return;
+        }
+        let source = sourcesByPort[port];
+        if (source) {
+          if (source.eventSource && source.url === url) return;
+
+          // How this has happened I don't understand...
+          // deregister from that source
+          const count = source.deregister(port);
+          // Clean-up
+          if (count === 0) {
+            source.close();
+            sourcesByUrl[source.url] = null;
+          }
+        }
+        // Create a new Source
+        source = new Source(url);
+        source.register(port);
+        sourcesByUrl[url] = source;
+        sourcesByPort[port] = source;
+      } else if (event.data.type === 'listen') {
+        const source = sourcesByPort[port];
+        source.listen(event.data.eventType);
+      } else if (event.data.type === 'close') {
+        const source = sourcesByPort[port];
+
+        if (!source) return;
+
+        const count = source.deregister(port);
+        if (count === 0) {
+          source.close();
+          sourcesByUrl[source.url] = null;
+          sourcesByPort[port] = null;
+        }
+      } else if (event.data.type === 'status') {
+        const source = sourcesByPort[port];
+        if (!source) {
+          port.postMessage({
+            type: 'status',
+            message: 'not connected',
+          });
+          return;
+        }
+        source.status(port);
+      } else {
+        // just send it back
+        port.postMessage({
+          type: 'error',
+          message: `received but don't know how to handle: ${event.data}`,
+        });
+      }
+    });
+    port.start();
+  }
+});

webpack.config.js

@@ -62,6 +62,9 @@ export default {
     'eventsource.sharedworker': [
       fileURLToPath(new URL('web_src/js/features/eventsource.sharedworker.js', import.meta.url)),
     ],
+    'websocket.sharedworker': [
+      fileURLToPath(new URL('web_src/js/features/websocket.sharedworker.js', import.meta.url)),
+    ],
     ...themes,
   },
   devtool: false,