Expose proxy to watch notifier WebSocket server in server mode: /.__marp-cli-watch-notifier__/*

Author: yhatt

src/converter.ts

@@ -41,7 +41,7 @@ import { isReadable } from './utils/finder'
 import { png2jpegViaPuppeteer } from './utils/jpeg'
 import { pdfLib, setOutline } from './utils/pdf'
 import { translateWSLPathToWindows } from './utils/wsl'
-import { notifier } from './watcher'
+import { notifier, type WatchNotifierEntrypointType } from './watcher'
 
 const CREATED_BY_MARP = 'Created by Marp'
 
@@ -93,7 +93,7 @@ export interface ConverterOption {
   templateOption?: TemplateOption
   themeSet: ThemeSet
   type: ConvertType
-  watch: boolean
+  watch: boolean | WatchNotifierEntrypointType
 }
 
 export interface ConvertFileOption {
@@ -184,7 +184,12 @@ export class Converter {
       base: await resolveBase(file),
       notifyWS:
         isFile(file) && this.options.watch && type === ConvertType.html
-          ? await notifier.register(file.absolutePath)
+          ? await notifier.register(
+              file.absolutePath,
+              typeof this.options.watch === 'string'
+                ? this.options.watch
+                : undefined
+            )
           : undefined,
       renderer: async (tplOpts) => {
         const engine = await this.generateEngine(tplOpts)

src/server.ts

@@ -20,6 +20,9 @@ import { CLIError, CLIErrorCode, error, isError } from './error'
 import { File, markdownExtensions } from './file'
 import serverIndex from './server/index.pug'
 import style from './server/index.scss'
+import { notifier } from './watcher'
+
+export const watchNotifierWebSocketEntrypoint = '.__marp-cli-watch-notifier__'
 
 export class Server extends (EventEmitter as new () => TypedEmitter<Server.Events>) {
   readonly converter: Converter
@@ -69,6 +72,20 @@ export class Server extends (EventEmitter as new () => TypedEmitter<Server.Event
           }
         })()
       )
+
+      this.httpServer.on('upgrade', (request, socket, head) => {
+        if (request.url?.startsWith(`/${watchNotifierWebSocketEntrypoint}/`)) {
+          const ws = notifier.server
+
+          if (ws) {
+            ws.handleUpgrade(request, socket, head, (client) => {
+              ws.emit('connection', client, request)
+            })
+            return
+          }
+        }
+        socket.destroy()
+      })
     })
   }
 
@@ -105,6 +122,7 @@ export class Server extends (EventEmitter as new () => TypedEmitter<Server.Event
     this.converter.options.output = false
     this.converter.options.pages = false
     this.converter.options.type = type
+    this.converter.options.watch = 'server'
 
     const result = await this.converter.convertFile(new File(filename))
     this.emit('converted', result)
@@ -179,6 +197,9 @@ export class Server extends (EventEmitter as new () => TypedEmitter<Server.Event
 
     this.server = express.default()
     this.server
+      .get(`/${watchNotifierWebSocketEntrypoint}/*all`, (_, res) => {
+        res.status(426).end('Upgrade Required')
+      })
       .get('*all', (req, res, next) =>
         this.preprocess(req, res).then(() => {
           if (!res.writableEnded) next()

src/utils/debug.ts

@@ -7,3 +7,4 @@ export const debugBrowserFinder = dbg('marp-cli:browser:finder')
 export const debugEngine = dbg('marp-cli:engine')
 export const debugPreview = dbg('marp-cli:preview')
 export const debugWatcher = dbg('marp-cli:watcher')
+export const debugWatcherWS = dbg('marp-cli:watcher:ws')

src/watcher.ts

@@ -8,7 +8,8 @@ import type { ServerOptions } from 'ws'
 import { Converter, ConvertedCallback } from './converter'
 import { isError } from './error'
 import { File, FileType } from './file'
-import { debugWatcher } from './utils/debug'
+import { watchNotifierWebSocketEntrypoint } from './server'
+import { debugWatcher, debugWatcherWS } from './utils/debug'
 
 const chokidarWatch: typeof _watch = (...args) => {
   debugWatcher('Start watching with chokidar: %O', args)
@@ -91,26 +92,47 @@ export class Watcher {
   }
 }
 
+export type WatchNotifierEntrypointType = 'static' | 'server'
+
 export class WatchNotifier {
   listeners = new Map<string, Set<any>>()
 
   private wss?: WebSocketServer
   private portNumber?: number
 
+  get server() {
+    return this.wss
+  }
+
   async port() {
     if (this.portNumber === undefined)
       this.portNumber = await getPortPromise({ port: 37717 })
 
     return this.portNumber
   }
 
-  async register(fn: string) {
+  async register(
+    fn: string,
+    entrypointType: WatchNotifierEntrypointType = 'static'
+  ) {
     const identifier = WatchNotifier.sha256(fn)
 
     if (!this.listeners.has(identifier))
       this.listeners.set(identifier, new Set())
 
-    return `ws://localhost:${await this.port()}/${identifier}`
+    return await this.entrypoint(identifier, entrypointType)
+  }
+
+  async entrypoint(
+    identifier: string,
+    entrypointType: WatchNotifierEntrypointType = 'static'
+  ) {
+    if (entrypointType === 'server') {
+      return `/${watchNotifierWebSocketEntrypoint}/${identifier}`
+    }
+
+    const port = await this.port()
+    return `ws://localhost:${port}/${identifier}`
   }
 
   sendTo(fn: string, command: string) {
@@ -124,21 +146,64 @@ export class WatchNotifier {
   }
 
   async start(opts: ServerOptions = {}) {
-    this.wss = new WebSocketServer({ ...opts, port: await this.port() })
-    this.wss.on('connection', (ws, sock) => {
-      if (sock.url) {
-        const [, identifier] = sock.url.split('/')
-        const wsSet = this.listeners.get(identifier)
+    const port = await this.port()
 
-        if (wsSet !== undefined) {
-          this.listeners.set(identifier, wsSet.add(ws))
+    this.wss = new WebSocketServer({ ...opts, port })
 
-          ws.on('close', () => this.listeners.get(identifier)!.delete(ws))
+    debugWatcherWS(
+      'WebSocket server for watch notifier started on port %d.',
+      port
+    )
 
-          ws.send('ready')
-          return
+    this.wss.on('connection', (ws, sock) => {
+      if (sock.url) {
+        debugWatcherWS('New WebSocket connection: %s', sock.url)
+
+        const identifier = (() => {
+          try {
+            const parsedUrl = new URL(sock.url, `ws://localhost:${port}`)
+            const detectedIdentifier = parsedUrl.pathname.split('/').pop()
+
+            debugWatcherWS(
+              'Detected identifier from WebSocket connection: %s',
+              detectedIdentifier
+            )
+
+            return detectedIdentifier
+          } catch (e: unknown) {
+            debugWatcherWS('Error occurred during parsing identifier: %o', e)
+            return undefined
+          }
+        })()
+
+        if (identifier) {
+          const wsSet = this.listeners.get(identifier)
+
+          if (wsSet !== undefined) {
+            this.listeners.set(identifier, wsSet.add(ws))
+            debugWatcherWS(
+              'WebSocket connection for identifier "%s" registered',
+              identifier
+            )
+
+            ws.on('close', () => {
+              this.listeners.get(identifier)!.delete(ws)
+              debugWatcherWS(
+                'WebSocket connection for identifier "%s" closed',
+                identifier
+              )
+            })
+
+            ws.send('ready')
+            return
+          }
         }
       }
+
+      debugWatcherWS(
+        'WebSocket connection request has been dismissed: %s',
+        sock.url
+      )
       ws.close()
     })
   }