From 9d4c60dccc8e7e777ad139a82e9f61feda9d21fc Mon Sep 17 00:00:00 2001
From: Chocobozzz <me@florianbigard.com>
Date: Tue, 11 Oct 2022 11:07:40 +0200
Subject: [PATCH] Add ability for plugins to register ws routes

---
 client/src/app/core/plugins/plugin.service.ts |  5 ++
 .../videos/shared/peertube-plugin.ts          |  1 +
 .../src/types/register-client-option.model.ts |  3 +
 server.ts                                     |  4 ++
 server/lib/plugins/plugin-helpers-builder.ts  | 13 ++--
 server/lib/plugins/plugin-manager.ts          | 31 +++++++-
 server/lib/plugins/register-helpers.ts        | 21 +++++-
 .../peertube-plugin-test-websocket/main.js    | 36 ++++++++++
 .../package.json                              | 20 ++++++
 server/tests/plugins/index.ts                 |  1 +
 server/tests/plugins/plugin-websocket.ts      | 70 +++++++++++++++++++
 server/types/plugins/index.ts                 |  1 +
 .../plugins/register-server-option.model.ts   | 14 ++++
 .../register-server-websocket-route.model.ts  |  8 +++
 support/doc/plugins/guide.md                  | 36 ++++++++++
 support/nginx/peertube                        |  5 ++
 16 files changed, 262 insertions(+), 7 deletions(-)
 create mode 100644 server/tests/fixtures/peertube-plugin-test-websocket/main.js
 create mode 100644 server/tests/fixtures/peertube-plugin-test-websocket/package.json
 create mode 100644 server/tests/plugins/plugin-websocket.ts
 create mode 100644 server/types/plugins/register-server-websocket-route.model.ts

diff --git a/client/src/app/core/plugins/plugin.service.ts b/client/src/app/core/plugins/plugin.service.ts
index dadc2a41d..1e79cbf79 100644
--- a/client/src/app/core/plugins/plugin.service.ts
+++ b/client/src/app/core/plugins/plugin.service.ts
@@ -202,6 +202,11 @@ export class PluginService implements ClientHook {
         return environment.apiUrl + `${pathPrefix}/${plugin.name}/${plugin.version}/router`
       },
 
+      getBaseWebSocketRoute: () => {
+        const pathPrefix = PluginsManager.getPluginPathPrefix(pluginInfo.isTheme)
+        return environment.apiUrl + `${pathPrefix}/${plugin.name}/${plugin.version}/ws`
+      },
+
       getBasePluginClientPath: () => {
         return '/p'
       },
diff --git a/client/src/standalone/videos/shared/peertube-plugin.ts b/client/src/standalone/videos/shared/peertube-plugin.ts
index 968854ce8..daf6f2b03 100644
--- a/client/src/standalone/videos/shared/peertube-plugin.ts
+++ b/client/src/standalone/videos/shared/peertube-plugin.ts
@@ -43,6 +43,7 @@ export class PeerTubePlugin {
     return {
       getBaseStaticRoute: unimplemented,
       getBaseRouterRoute: unimplemented,
+      getBaseWebSocketRoute: unimplemented,
       getBasePluginClientPath: unimplemented,
 
       getSettings: () => {
diff --git a/client/src/types/register-client-option.model.ts b/client/src/types/register-client-option.model.ts
index 2460a7499..2c09f15a7 100644
--- a/client/src/types/register-client-option.model.ts
+++ b/client/src/types/register-client-option.model.ts
@@ -24,6 +24,9 @@ export type RegisterClientHelpers = {
 
   getBaseRouterRoute: () => string
 
+  // PeerTube >= 5.0
+  getBaseWebSocketRoute: () => string
+
   getBasePluginClientPath: () => string
 
   isLoggedIn: () => boolean
diff --git a/server.ts b/server.ts
index 417387a4f..a29b5e408 100644
--- a/server.ts
+++ b/server.ts
@@ -328,6 +328,10 @@ async function startApplication () {
   GeoIPUpdateScheduler.Instance.enable()
   OpenTelemetryMetrics.Instance.registerMetrics()
 
+  PluginManager.Instance.init(server)
+  // Before PeerTubeSocket init
+  PluginManager.Instance.registerWebSocketRouter()
+
   PeerTubeSocket.Instance.init(server)
   VideoViewsManager.Instance.init()
 
diff --git a/server/lib/plugins/plugin-helpers-builder.ts b/server/lib/plugins/plugin-helpers-builder.ts
index 35945422c..7b1def6e3 100644
--- a/server/lib/plugins/plugin-helpers-builder.ts
+++ b/server/lib/plugins/plugin-helpers-builder.ts
@@ -1,4 +1,5 @@
 import express from 'express'
+import { Server } from 'http'
 import { join } from 'path'
 import { ffprobePromise } from '@server/helpers/ffmpeg/ffprobe-utils'
 import { buildLogger } from '@server/helpers/logger'
@@ -17,12 +18,12 @@ import { MPlugin, MVideo, UserNotificationModelForApi } from '@server/types/mode
 import { PeerTubeHelpers } from '@server/types/plugins'
 import { VideoBlacklistCreate, VideoStorage } from '@shared/models'
 import { addAccountInBlocklist, addServerInBlocklist, removeAccountFromBlocklist, removeServerFromBlocklist } from '../blocklist'
+import { PeerTubeSocket } from '../peertube-socket'
 import { ServerConfigManager } from '../server-config-manager'
 import { blacklistVideo, unblacklistVideo } from '../video-blacklist'
 import { VideoPathManager } from '../video-path-manager'
-import { PeerTubeSocket } from '../peertube-socket'
 
-function buildPluginHelpers (pluginModel: MPlugin, npmName: string): PeerTubeHelpers {
+function buildPluginHelpers (httpServer: Server, pluginModel: MPlugin, npmName: string): PeerTubeHelpers {
   const logger = buildPluginLogger(npmName)
 
   const database = buildDatabaseHelpers()
@@ -30,7 +31,7 @@ function buildPluginHelpers (pluginModel: MPlugin, npmName: string): PeerTubeHel
 
   const config = buildConfigHelpers()
 
-  const server = buildServerHelpers()
+  const server = buildServerHelpers(httpServer)
 
   const moderation = buildModerationHelpers()
 
@@ -69,8 +70,10 @@ function buildDatabaseHelpers () {
   }
 }
 
-function buildServerHelpers () {
+function buildServerHelpers (httpServer: Server) {
   return {
+    getHTTPServer: () => httpServer,
+
     getServerActor: () => getServerActor()
   }
 }
@@ -218,6 +221,8 @@ function buildPluginRelatedHelpers (plugin: MPlugin, npmName: string) {
 
     getBaseRouterRoute: () => `/plugins/${plugin.name}/${plugin.version}/router/`,
 
+    getBaseWebSocketRoute: () => `/plugins/${plugin.name}/${plugin.version}/ws/`,
+
     getDataDirectoryPath: () => join(CONFIG.STORAGE.PLUGINS_DIR, 'data', npmName)
   }
 }
diff --git a/server/lib/plugins/plugin-manager.ts b/server/lib/plugins/plugin-manager.ts
index a46b97fa4..c4d9b6574 100644
--- a/server/lib/plugins/plugin-manager.ts
+++ b/server/lib/plugins/plugin-manager.ts
@@ -1,6 +1,7 @@
 import express from 'express'
 import { createReadStream, createWriteStream } from 'fs'
 import { ensureDir, outputFile, readJSON } from 'fs-extra'
+import { Server } from 'http'
 import { basename, join } from 'path'
 import { decachePlugin } from '@server/helpers/decache'
 import { ApplicationModel } from '@server/models/application/application'
@@ -67,9 +68,37 @@ export class PluginManager implements ServerHook {
   private hooks: { [name: string]: HookInformationValue[] } = {}
   private translations: PluginLocalesTranslations = {}
 
+  private server: Server
+
   private constructor () {
   }
 
+  init (server: Server) {
+    this.server = server
+  }
+
+  registerWebSocketRouter () {
+    this.server.on('upgrade', (request, socket, head) => {
+      const url = request.url
+
+      const matched = url.match(`/plugins/([^/]+)/([^/]+/)?ws(/.*)`)
+      if (!matched) return
+
+      const npmName = PluginModel.buildNpmName(matched[1], PluginType.PLUGIN)
+      const subRoute = matched[3]
+
+      const result = this.getRegisteredPluginOrTheme(npmName)
+      if (!result) return
+
+      const routes = result.registerHelpers.getWebSocketRoutes()
+
+      const wss = routes.find(r => r.route.startsWith(subRoute))
+      if (!wss) return
+
+      wss.handler(request, socket, head)
+    })
+  }
+
   // ###################### Getters ######################
 
   isRegistered (npmName: string) {
@@ -581,7 +610,7 @@ export class PluginManager implements ServerHook {
       })
     }
 
-    const registerHelpers = new RegisterHelpers(npmName, plugin, onHookAdded.bind(this))
+    const registerHelpers = new RegisterHelpers(npmName, plugin, this.server, onHookAdded.bind(this))
 
     return {
       registerStore: registerHelpers,
diff --git a/server/lib/plugins/register-helpers.ts b/server/lib/plugins/register-helpers.ts
index f4d405676..1aaef3606 100644
--- a/server/lib/plugins/register-helpers.ts
+++ b/server/lib/plugins/register-helpers.ts
@@ -1,4 +1,5 @@
 import express from 'express'
+import { Server } from 'http'
 import { logger } from '@server/helpers/logger'
 import { onExternalUserAuthenticated } from '@server/lib/auth/external-auth'
 import { VideoConstantManagerFactory } from '@server/lib/plugins/video-constant-manager-factory'
@@ -8,7 +9,8 @@ import {
   RegisterServerAuthExternalResult,
   RegisterServerAuthPassOptions,
   RegisterServerExternalAuthenticatedResult,
-  RegisterServerOptions
+  RegisterServerOptions,
+  RegisterServerWebSocketRouteOptions
 } from '@server/types/plugins'
 import {
   EncoderOptionsBuilder,
@@ -49,12 +51,15 @@ export class RegisterHelpers {
 
   private readonly onSettingsChangeCallbacks: SettingsChangeCallback[] = []
 
+  private readonly webSocketRoutes: RegisterServerWebSocketRouteOptions[] = []
+
   private readonly router: express.Router
   private readonly videoConstantManagerFactory: VideoConstantManagerFactory
 
   constructor (
     private readonly npmName: string,
     private readonly plugin: PluginModel,
+    private readonly server: Server,
     private readonly onHookAdded: (options: RegisterServerHookOptions) => void
   ) {
     this.router = express.Router()
@@ -66,6 +71,7 @@ export class RegisterHelpers {
     const registerSetting = this.buildRegisterSetting()
 
     const getRouter = this.buildGetRouter()
+    const registerWebSocketRoute = this.buildRegisterWebSocketRoute()
 
     const settingsManager = this.buildSettingsManager()
     const storageManager = this.buildStorageManager()
@@ -85,13 +91,14 @@ export class RegisterHelpers {
     const unregisterIdAndPassAuth = this.buildUnregisterIdAndPassAuth()
     const unregisterExternalAuth = this.buildUnregisterExternalAuth()
 
-    const peertubeHelpers = buildPluginHelpers(this.plugin, this.npmName)
+    const peertubeHelpers = buildPluginHelpers(this.server, this.plugin, this.npmName)
 
     return {
       registerHook,
       registerSetting,
 
       getRouter,
+      registerWebSocketRoute,
 
       settingsManager,
       storageManager,
@@ -180,10 +187,20 @@ export class RegisterHelpers {
     return this.onSettingsChangeCallbacks
   }
 
+  getWebSocketRoutes () {
+    return this.webSocketRoutes
+  }
+
   private buildGetRouter () {
     return () => this.router
   }
 
+  private buildRegisterWebSocketRoute () {
+    return (options: RegisterServerWebSocketRouteOptions) => {
+      this.webSocketRoutes.push(options)
+    }
+  }
+
   private buildRegisterSetting () {
     return (options: RegisterServerSettingOptions) => {
       this.settings.push(options)
diff --git a/server/tests/fixtures/peertube-plugin-test-websocket/main.js b/server/tests/fixtures/peertube-plugin-test-websocket/main.js
new file mode 100644
index 000000000..3fde76cfe
--- /dev/null
+++ b/server/tests/fixtures/peertube-plugin-test-websocket/main.js
@@ -0,0 +1,36 @@
+const WebSocketServer = require('ws').WebSocketServer
+
+async function register ({
+  registerWebSocketRoute
+}) {
+  const wss = new WebSocketServer({ noServer: true })
+
+  wss.on('connection', function connection(ws) {
+    ws.on('message', function message(data) {
+      if (data.toString() === 'ping') {
+        ws.send('pong')
+      }
+    })
+  })
+
+  registerWebSocketRoute({
+    route: '/toto',
+
+    handler: (request, socket, head) => {
+      wss.handleUpgrade(request, socket, head, ws => {
+        wss.emit('connection', ws, request)
+      })
+    }
+  })
+}
+
+async function unregister () {
+  return
+}
+
+module.exports = {
+  register,
+  unregister
+}
+
+// ###########################################################################
diff --git a/server/tests/fixtures/peertube-plugin-test-websocket/package.json b/server/tests/fixtures/peertube-plugin-test-websocket/package.json
new file mode 100644
index 000000000..89c8baa04
--- /dev/null
+++ b/server/tests/fixtures/peertube-plugin-test-websocket/package.json
@@ -0,0 +1,20 @@
+{
+  "name": "peertube-plugin-test-websocket",
+  "version": "0.0.1",
+  "description": "Plugin test websocket",
+  "engine": {
+    "peertube": ">=1.3.0"
+  },
+  "keywords": [
+    "peertube",
+    "plugin"
+  ],
+  "homepage": "https://github.com/Chocobozzz/PeerTube",
+  "author": "Chocobozzz",
+  "bugs": "https://github.com/Chocobozzz/PeerTube/issues",
+  "library": "./main.js",
+  "staticDirs": {},
+  "css": [],
+  "clientScripts": [],
+  "translations": {}
+}
diff --git a/server/tests/plugins/index.ts b/server/tests/plugins/index.ts
index 4534120fd..210af7236 100644
--- a/server/tests/plugins/index.ts
+++ b/server/tests/plugins/index.ts
@@ -8,5 +8,6 @@ import './plugin-router'
 import './plugin-storage'
 import './plugin-transcoding'
 import './plugin-unloading'
+import './plugin-websocket'
 import './translations'
 import './video-constants'
diff --git a/server/tests/plugins/plugin-websocket.ts b/server/tests/plugins/plugin-websocket.ts
new file mode 100644
index 000000000..adaa28b1d
--- /dev/null
+++ b/server/tests/plugins/plugin-websocket.ts
@@ -0,0 +1,70 @@
+/* eslint-disable @typescript-eslint/no-unused-expressions,@typescript-eslint/require-await */
+
+import WebSocket from 'ws'
+import { cleanupTests, createSingleServer, PeerTubeServer, PluginsCommand, setAccessTokensToServers } from '@shared/server-commands'
+
+function buildWebSocket (server: PeerTubeServer, path: string) {
+  return new WebSocket('ws://' + server.host + path)
+}
+
+function expectErrorOrTimeout (server: PeerTubeServer, path: string, expectedTimeout: number) {
+  return new Promise<void>((res, rej) => {
+    const ws = buildWebSocket(server, path)
+    ws.on('error', () => res())
+
+    const timeout = setTimeout(() => res(), expectedTimeout)
+
+    ws.on('open', () => {
+      clearTimeout(timeout)
+
+      return rej(new Error('Connect did not timeout'))
+    })
+  })
+}
+
+describe('Test plugin websocket', function () {
+  let server: PeerTubeServer
+  const basePaths = [
+    '/plugins/test-websocket/ws/',
+    '/plugins/test-websocket/0.0.1/ws/'
+  ]
+
+  before(async function () {
+    this.timeout(30000)
+
+    server = await createSingleServer(1)
+    await setAccessTokensToServers([ server ])
+
+    await server.plugins.install({ path: PluginsCommand.getPluginTestPath('-websocket') })
+  })
+
+  it('Should not connect to the websocket without the appropriate path', async function () {
+    const paths = [
+      '/plugins/unknown/ws/',
+      '/plugins/unknown/0.0.1/ws/'
+    ]
+
+    for (const path of paths) {
+      await expectErrorOrTimeout(server, path, 1000)
+    }
+  })
+
+  it('Should not connect to the websocket without the appropriate sub path', async function () {
+    for (const path of basePaths) {
+      await expectErrorOrTimeout(server, path + '/unknown', 1000)
+    }
+  })
+
+  it('Should connect to the websocket and receive pong', function (done) {
+    const ws = buildWebSocket(server, basePaths[0])
+
+    ws.on('open', () => ws.send('ping'))
+    ws.on('message', data => {
+      if (data.toString() === 'pong') return done()
+    })
+  })
+
+  after(async function () {
+    await cleanupTests([ server ])
+  })
+})
diff --git a/server/types/plugins/index.ts b/server/types/plugins/index.ts
index de30ff2ab..bf9c35d49 100644
--- a/server/types/plugins/index.ts
+++ b/server/types/plugins/index.ts
@@ -1,3 +1,4 @@
 export * from './plugin-library.model'
 export * from './register-server-auth.model'
 export * from './register-server-option.model'
+export * from './register-server-websocket-route.model'
diff --git a/server/types/plugins/register-server-option.model.ts b/server/types/plugins/register-server-option.model.ts
index a8b804b63..1e2bd830e 100644
--- a/server/types/plugins/register-server-option.model.ts
+++ b/server/types/plugins/register-server-option.model.ts
@@ -1,4 +1,5 @@
 import { Response, Router } from 'express'
+import { Server } from 'http'
 import { Logger } from 'winston'
 import { ActorModel } from '@server/models/actor/actor'
 import {
@@ -22,6 +23,7 @@ import {
   RegisterServerAuthExternalResult,
   RegisterServerAuthPassOptions
 } from './register-server-auth.model'
+import { RegisterServerWebSocketRouteOptions } from './register-server-websocket-route.model'
 
 export type PeerTubeHelpers = {
   logger: Logger
@@ -83,6 +85,9 @@ export type PeerTubeHelpers = {
   }
 
   server: {
+    // PeerTube >= 5.0
+    getHTTPServer: () => Server
+
     getServerActor: () => Promise<ActorModel>
   }
 
@@ -97,6 +102,8 @@ export type PeerTubeHelpers = {
 
     // PeerTube >= 3.2
     getBaseRouterRoute: () => string
+    // PeerTube >= 5.0
+    getBaseWebSocketRoute: () => string
 
     // PeerTube >= 3.2
     getDataDirectoryPath: () => string
@@ -140,5 +147,12 @@ export type RegisterServerOptions = {
   //  * /plugins/:pluginName/router/...
   getRouter(): Router
 
+  // PeerTube >= 5.0
+  // Register WebSocket route
+  // Base routes of the WebSocket router are
+  //  * /plugins/:pluginName/:pluginVersion/ws/...
+  //  * /plugins/:pluginName/ws/...
+  registerWebSocketRoute: (options: RegisterServerWebSocketRouteOptions) => void
+
   peertubeHelpers: PeerTubeHelpers
 }
diff --git a/server/types/plugins/register-server-websocket-route.model.ts b/server/types/plugins/register-server-websocket-route.model.ts
new file mode 100644
index 000000000..edf64f66b
--- /dev/null
+++ b/server/types/plugins/register-server-websocket-route.model.ts
@@ -0,0 +1,8 @@
+import { IncomingMessage } from 'http'
+import { Duplex } from 'stream'
+
+export type RegisterServerWebSocketRouteOptions = {
+  route: string
+
+  handler: (request: IncomingMessage, socket: Duplex, head: Buffer) => any
+}
diff --git a/support/doc/plugins/guide.md b/support/doc/plugins/guide.md
index 431d5332f..1c809258a 100644
--- a/support/doc/plugins/guide.md
+++ b/support/doc/plugins/guide.md
@@ -12,6 +12,7 @@
     - [Storage](#storage)
     - [Update video constants](#update-video-constants)
     - [Add custom routes](#add-custom-routes)
+    - [Add custom WebSocket handlers](#add-custom-websocket-handlers)
     - [Add external auth methods](#add-external-auth-methods)
     - [Add new transcoding profiles](#add-new-transcoding-profiles)
     - [Server helpers](#server-helpers)
@@ -317,6 +318,41 @@ The `ping` route can be accessed using:
  * Or `/plugins/:pluginName/router/ping`
 
 
+#### Add custom WebSocket handlers
+
+You can create custom WebSocket servers (like [ws](https://github.com/websockets/ws) for example) using `registerWebSocketRoute`:
+
+```js
+function register ({
+  registerWebSocketRoute,
+  peertubeHelpers
+}) {
+  const wss = new WebSocketServer({ noServer: true })
+
+  wss.on('connection', function connection(ws) {
+    peertubeHelpers.logger.info('WebSocket connected!')
+
+    setInterval(() => {
+      ws.send('WebSocket message sent by server');
+    }, 1000)
+  })
+
+  registerWebSocketRoute({
+    route: '/my-websocket-route',
+
+    handler: (request, socket, head) => {
+      wss.handleUpgrade(request, socket, head, ws => {
+        wss.emit('connection', ws, request)
+      })
+    }
+  })
+}
+```
+
+The `my-websocket-route` route can be accessed using:
+ * `/plugins/:pluginName/:pluginVersion/ws/my-websocket-route`
+ * Or `/plugins/:pluginName/ws/my-websocket-route`
+
 #### Add external auth methods
 
 If you want to add a classic username/email and password auth method (like [LDAP](https://framagit.org/framasoft/peertube/official-plugins/-/tree/master/peertube-plugin-auth-ldap) for example):
diff --git a/support/nginx/peertube b/support/nginx/peertube
index abb83d5c4..f6f754b58 100644
--- a/support/nginx/peertube
+++ b/support/nginx/peertube
@@ -132,6 +132,11 @@ server {
     try_files /dev/null @api_websocket;
   }
 
+  # Plugin websocket routes
+  location ~ ^/plugins/[^/]+(/[^/]+)?/ws/ {
+    try_files /dev/null @api_websocket;
+  }
+
   ##
   # Performance optimizations
   # For extra performance please refer to https://github.com/denji/nginx-tuning
-- 
2.41.0