yusukebe · May 21, 2025 02:12
diff --git a/plugin.ts b/plugin.ts
 import { Plugin, normalizePath } from 'vite'
 import path from 'node:path'
 import picomatch from 'picomatch'
 import type { ServerResponse } from 'node:http'

 type ShouldInjectFunction = (req: ServerResponse['req'], res: ServerResponse) => boolean

 type Options = {
  /**
   * default ['src\/\*\*\/\*.ts', 'src\/\*\*\/\*.tsx']
   */
  entry?: string | string[]

  ignore?: string | string[]
  /**
   * default: true
   */
  injectViteClient?: boolean | ShouldInjectFunction
  /**
   * default: false
   */
  injectReactRefresh?: boolean | ShouldInjectFunction
 }

 function intoShouldInjectFunction(
  value: boolean | ShouldInjectFunction | undefined,
  defaultValue: boolean
 ): ShouldInjectFunction {
  switch (typeof value) {
    case 'boolean':
      return (_req, _res) => value
    case 'function':
      return value
    default:
      return (_req, _res) => defaultValue
  }
 }

 // HTMLにスクリプトを挿入する関数
 function injectScripts(html: string, shouldInjectViteClient: boolean, shouldInjectReactRefresh: boolean): string {
  const hasRefresh = html.includes('/@react-refresh')
  const hasViteClient = html.includes('/@vite/client')

  let injection = ''

  if (shouldInjectReactRefresh && !hasRefresh) {
    injection += `<script type="module" src="/@react-refresh"></script>
 <script type="module">
  import RefreshRuntime from '/@react-refresh'
  RefreshRuntime.injectIntoGlobalHook(window)
  window.$RefreshReg$ = () => {}
  window.$RefreshSig$ = () => (type) => type
  window.__vite_plugin_react_preamble_installed__ = true
 </script>\n`
  }

  if (shouldInjectViteClient && !hasViteClient) {
    injection += `<script type="module" src="/@vite/client"></script>\n`
  }

  if (!injection) {
    return html
  }

  let result = html

  if (result.includes('<head>')) {
    result = result.replace('<head>', `<head>\n${injection}`)
  } else if (result.includes('<html>')) {
    result = result.replace('<html>', `<html>\n${injection}`)
  } else {
    result = `${injection}${result}`
  }

  return result
 }

 export default function ssrHotReload(options: Options = {}): Plugin {
  const entryPatterns = Array.isArray(options.entry)
    ? options.entry
    : options.entry
    ? [options.entry]
    : ['src/**/*.ts', 'src/**/*.tsx']

  const ignorePatterns = Array.isArray(options.ignore) ? options.ignore : options.ignore ? [options.ignore] : []

  const injectReactRefresh = intoShouldInjectFunction(options.injectReactRefresh, false)
  const injectViteClientFn = intoShouldInjectFunction(options.injectViteClient, true)

  let root = process.cwd()
  let isMatch: (file: string) => boolean

  return {
    name: 'vite-plugin-ssr-hot-reload',
    apply: 'serve',

    configResolved(config) {
      root = config.root || process.cwd()

      const normalizedEntries = entryPatterns.map((p) => normalizeGlobPattern(p, root))
      const normalizedIgnores = ignorePatterns.map((p) => normalizeGlobPattern(p, root))

      const matcher = picomatch(normalizedEntries, {
        ignore: normalizedIgnores,
        dot: true
      })

      isMatch = (filePath: string) => {
        const rel = normalizePath(path.relative(root, filePath))
        return matcher(rel)
      }
    },

    configureServer(server) {
      // HTMLレスポンスを処理する早期ミドルウェア
      server.middlewares.use(function (req, res, next) {
        // 最初に、これがHTMLレスポンスになる可能性があるかどうかを確認
        const url = req.url || ''
        const acceptHeader = req.headers.accept || ''

        // HTML以外のリクエストはそのまま通過させる（効率化のため）
        if (url.match(/\.(js|css|png|jpg|jpeg|gif|svg|ico|json|map)(\?.*)?$/) || !acceptHeader.includes('text/html')) {
          return next()
        }

        // レスポンスのwriteHeadメソッドをオーバーライド
        const originalWriteHead = res.writeHead
        res.writeHead = function (statusCode: number, ...args: any[]) {
          // HTMLコンテンツかどうかを確認
          const headers = args[0] && typeof args[0] !== 'string' ? args[0] : args[1]
          let contentType = ''

          if (headers) {
            Object.keys(headers).forEach((key) => {
              if (key.toLowerCase() === 'content-type') {
                contentType = headers[key]
              }
            })
          }

          // レスポンスヘッダーからContent-Typeを取得
          if (!contentType) {
            contentType = typeof res.getHeader === 'function' ? (res.getHeader('content-type') as string) || '' : ''
          }

          // HTMLレスポンスでない場合、元の動作を続行
          if (!contentType || !contentType.includes('text/html')) {
            return originalWriteHead.apply(res, [statusCode, ...args])
          }

          // ここからHTMLレスポンスの処理
          let _resolve: (value: any) => void
          let _reject: (reason?: any) => void
          const collectorPromise = new Promise<Buffer[]>((resolve, reject) => {
            _resolve = resolve
            _reject = reject
          })

          // レスポンスデータを収集するための配列
          const chunks: Buffer[] = []

          // 元のメソッドを保存
          const originalWrite = res.write
          const originalEnd = res.end

          // writeメソッドをオーバーライド
          res.write = function (chunk: any, ...args: any[]) {
            if (chunk) {
              const buffer = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk)
              chunks.push(buffer)
            }
            return true // ストリームが続行中
          }

          // endメソッドをオーバーライド
          res.end = function (chunk: any, ...args: any[]) {
            if (chunk) {
              const buffer = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk)
              chunks.push(buffer)
            }

            // 収集完了を通知
            _resolve(chunks)
            return true // ストリームが終了
          }

          // 非同期でレスポンス処理を実行
          collectorPromise
            .then((buffers) => {
              // 元のHTMLを取得
              const html = Buffer.concat(buffers).toString()

              // スクリプトを挿入
              const modifiedHtml = injectScripts(html, injectViteClientFn(req, res), injectReactRefresh(req, res))

              // 元のwriteとendメソッドを復元
              res.write = originalWrite
              res.end = originalEnd

              // コンテンツ長によってヘッダーを設定
              if (headers && typeof headers === 'object') {
                // 既存のヘッダーからContent-Lengthを削除
                if (headers['content-length'] || headers['Content-Length']) {
                  delete headers['content-length']
                  delete headers['Content-Length']
                }

                // Transfer-Encodingをchunkedに設定
                headers['transfer-encoding'] = 'chunked'
              } else {
                // writtenHeadersがない場合は新しいヘッダーオブジェクトを作成
                args[0] = { 'transfer-encoding': 'chunked' }
              }

              // ステータスコードとヘッダーを書き込む
              originalWriteHead.apply(res, [statusCode, ...args])

              // 変更されたHTMLを書き込む
              originalWrite.call(res, modifiedHtml)
              originalEnd.call(res)
            })
            .catch((err) => {
              console.error('Error in HTML processing:', err)
              // エラーの場合は元のメソッドを復元して続行
              res.write = originalWrite
              res.end = originalEnd
              next(err)
            })

          // オリジナルのwriteHeadをスキップ（後で呼び出す）
          return true
        }

        next()
      })
    },

    handleHotUpdate({ server, file }) {
      if (!file) return

      if (isMatch?.(file)) {
        server.hot.send({ type: 'full-reload' })
        return []
      }
    }
  }

  function normalizeGlobPattern(pattern: string, root: string): string {
    const normalized = normalizePath(pattern)

    if (path.isAbsolute(normalized)) {
      const relative = path.relative(root, normalized)
      if (!relative.startsWith('..') && !path.isAbsolute(relative)) {
        return normalizePath(relative)
      }
      return normalized.slice(1)
    }

    if (normalized.startsWith('/')) {
      return normalized.slice(1)
    }

    if (normalized.startsWith('./')) {
      return normalized.slice(2)
    }

    return normalized
  }
 }
	import { Plugin, normalizePath } from 'vite'
	import path from 'node:path'
	import picomatch from 'picomatch'
	import type { ServerResponse } from 'node:http'

	type ShouldInjectFunction = (req: ServerResponse['req'], res: ServerResponse) => boolean

	type Options = {
	/**
	* default ['src\/\\\/\.ts', 'src\/\\\/\.tsx']
	*/
	entry?: string \| string[]

	ignore?: string \| string[]
	/**
	* default: true
	*/
	injectViteClient?: boolean \| ShouldInjectFunction
	/**
	* default: false
	*/
	injectReactRefresh?: boolean \| ShouldInjectFunction
	}

	function intoShouldInjectFunction(
	value: boolean \| ShouldInjectFunction \| undefined,
	defaultValue: boolean
	): ShouldInjectFunction {
	switch (typeof value) {
	case 'boolean':
	return (_req, _res) => value
	case 'function':
	return value
	default:
	return (_req, _res) => defaultValue
	}
	}

	// HTMLにスクリプトを挿入する関数
	function injectScripts(html: string, shouldInjectViteClient: boolean, shouldInjectReactRefresh: boolean): string {
	const hasRefresh = html.includes('/@react-refresh')
	const hasViteClient = html.includes('/@vite/client')

	let injection = ''

	if (shouldInjectReactRefresh && !hasRefresh) {
	injection += `<script type="module" src="/@react-refresh"></script>
	<script type="module">
	import RefreshRuntime from '/@react-refresh'
	RefreshRuntime.injectIntoGlobalHook(window)
	window.$RefreshReg$ = () => {}
	window.$RefreshSig$ = () => (type) => type
	window.__vite_plugin_react_preamble_installed__ = true
	</script>\n`
	}

	if (shouldInjectViteClient && !hasViteClient) {
	injection += `<script type="module" src="/@vite/client"></script>\n`
	}

	if (!injection) {
	return html
	}

	let result = html

	if (result.includes('<head>')) {
	result = result.replace('<head>', `<head>\n${injection}`)
	} else if (result.includes('<html>')) {
	result = result.replace('<html>', `<html>\n${injection}`)
	} else {
	result = `${injection}${result}`
	}

	return result
	}

	export default function ssrHotReload(options: Options = {}): Plugin {
	const entryPatterns = Array.isArray(options.entry)
	? options.entry
	: options.entry
	? [options.entry]
	: ['src/*/.ts', 'src/*/.tsx']

	const ignorePatterns = Array.isArray(options.ignore) ? options.ignore : options.ignore ? [options.ignore] : []

	const injectReactRefresh = intoShouldInjectFunction(options.injectReactRefresh, false)
	const injectViteClientFn = intoShouldInjectFunction(options.injectViteClient, true)

	let root = process.cwd()
	let isMatch: (file: string) => boolean

	return {
	name: 'vite-plugin-ssr-hot-reload',
	apply: 'serve',

	configResolved(config) {
	root = config.root \|\| process.cwd()

	const normalizedEntries = entryPatterns.map((p) => normalizeGlobPattern(p, root))
	const normalizedIgnores = ignorePatterns.map((p) => normalizeGlobPattern(p, root))

	const matcher = picomatch(normalizedEntries, {
	ignore: normalizedIgnores,
	dot: true
	})

	isMatch = (filePath: string) => {
	const rel = normalizePath(path.relative(root, filePath))
	return matcher(rel)
	}
	},

	configureServer(server) {
	// HTMLレスポンスを処理する早期ミドルウェア
	server.middlewares.use(function (req, res, next) {
	// 最初に、これがHTMLレスポンスになる可能性があるかどうかを確認
	const url = req.url \|\| ''
	const acceptHeader = req.headers.accept \|\| ''

	// HTML以外のリクエストはそのまま通過させる（効率化のため）
	if (url.match(/\.(js\|css\|png\|jpg\|jpeg\|gif\|svg\|ico\|json\|map)(\?.*)?$/) \|\| !acceptHeader.includes('text/html')) {
	return next()
	}

	// レスポンスのwriteHeadメソッドをオーバーライド
	const originalWriteHead = res.writeHead
	res.writeHead = function (statusCode: number, ...args: any[]) {
	// HTMLコンテンツかどうかを確認
	const headers = args[0] && typeof args[0] !== 'string' ? args[0] : args[1]
	let contentType = ''

	if (headers) {
	Object.keys(headers).forEach((key) => {
	if (key.toLowerCase() === 'content-type') {
	contentType = headers[key]
	}
	})
	}

	// レスポンスヘッダーからContent-Typeを取得
	if (!contentType) {
	contentType = typeof res.getHeader === 'function' ? (res.getHeader('content-type') as string) \|\| '' : ''
	}

	// HTMLレスポンスでない場合、元の動作を続行
	if (!contentType \|\| !contentType.includes('text/html')) {
	return originalWriteHead.apply(res, [statusCode, ...args])
	}

	// ここからHTMLレスポンスの処理
	let _resolve: (value: any) => void
	let _reject: (reason?: any) => void
	const collectorPromise = new Promise<Buffer[]>((resolve, reject) => {
	_resolve = resolve
	_reject = reject
	})

	// レスポンスデータを収集するための配列
	const chunks: Buffer[] = []

	// 元のメソッドを保存
	const originalWrite = res.write
	const originalEnd = res.end

	// writeメソッドをオーバーライド
	res.write = function (chunk: any, ...args: any[]) {
	if (chunk) {
	const buffer = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk)
	chunks.push(buffer)
	}
	return true // ストリームが続行中
	}

	// endメソッドをオーバーライド
	res.end = function (chunk: any, ...args: any[]) {
	if (chunk) {
	const buffer = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk)
	chunks.push(buffer)
	}

	// 収集完了を通知
	_resolve(chunks)
	return true // ストリームが終了
	}

	// 非同期でレスポンス処理を実行
	collectorPromise
	.then((buffers) => {
	// 元のHTMLを取得
	const html = Buffer.concat(buffers).toString()

	// スクリプトを挿入
	const modifiedHtml = injectScripts(html, injectViteClientFn(req, res), injectReactRefresh(req, res))

	// 元のwriteとendメソッドを復元
	res.write = originalWrite
	res.end = originalEnd

	// コンテンツ長によってヘッダーを設定
	if (headers && typeof headers === 'object') {
	// 既存のヘッダーからContent-Lengthを削除
	if (headers['content-length'] \|\| headers['Content-Length']) {
	delete headers['content-length']
	delete headers['Content-Length']
	}

	// Transfer-Encodingをchunkedに設定
	headers['transfer-encoding'] = 'chunked'
	} else {
	// writtenHeadersがない場合は新しいヘッダーオブジェクトを作成
	args[0] = { 'transfer-encoding': 'chunked' }
	}

	// ステータスコードとヘッダーを書き込む
	originalWriteHead.apply(res, [statusCode, ...args])

	// 変更されたHTMLを書き込む
	originalWrite.call(res, modifiedHtml)
	originalEnd.call(res)
	})
	.catch((err) => {
	console.error('Error in HTML processing:', err)
	// エラーの場合は元のメソッドを復元して続行
	res.write = originalWrite
	res.end = originalEnd
	next(err)
	})

	// オリジナルのwriteHeadをスキップ（後で呼び出す）
	return true
	}

	next()
	})
	},

	handleHotUpdate({ server, file }) {
	if (!file) return

	if (isMatch?.(file)) {
	server.hot.send({ type: 'full-reload' })
	return []
	}
	}
	}

	function normalizeGlobPattern(pattern: string, root: string): string {
	const normalized = normalizePath(pattern)

	if (path.isAbsolute(normalized)) {
	const relative = path.relative(root, normalized)
	if (!relative.startsWith('..') && !path.isAbsolute(relative)) {
	return normalizePath(relative)
	}
	return normalized.slice(1)
	}

	if (normalized.startsWith('/')) {
	return normalized.slice(1)
	}

	if (normalized.startsWith('./')) {
	return normalized.slice(2)
	}

	return normalized
	}
	}