Add ability for plugins to register ws routes

10 files changed, 208 insertions, 7 deletions

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
+}