From b6e91e06ed455cd11a7e22927bad348d63556f31 Mon Sep 17 00:00:00 2001
From: Evan You <yyx990803@gmail.com>
Date: Thu, 2 Jul 2020 12:48:59 -0400
Subject: [PATCH] feat: connect to actual vite server port for websocket

Allows assets to be served under a different port
close #452, close #460
---
 src/client/client.ts               |  3 ++-
 src/node/server/index.ts           | 11 ++++++++---
 src/node/server/serverPluginHmr.ts |  2 +-
 3 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/src/client/client.ts b/src/client/client.ts
index aa9a3e55bbafbf..5448bf703e7386 100644
--- a/src/client/client.ts
+++ b/src/client/client.ts
@@ -1,5 +1,6 @@
 // replaced by serverPluginHmr when served
 declare const __SW_ENABLED__: boolean
+declare const __PORT__: number
 
 // This file runs in the browser.
 import { HMRRuntime } from 'vue'
@@ -52,7 +53,7 @@ console.log('[vite] connecting...')
 declare var __VUE_HMR_RUNTIME__: HMRRuntime
 
 const socketProtocol = location.protocol === 'https:' ? 'wss' : 'ws'
-const socketUrl = `${socketProtocol}://${location.host}`
+const socketUrl = `${socketProtocol}://${location.hostname}:${__PORT__}`
 const socket = new WebSocket(socketUrl, 'vite-hmr')
 
 function warnFailedFetch(err: Error, path: string | string[]) {
diff --git a/src/node/server/index.ts b/src/node/server/index.ts
index 2512b449679d95..432c86f2965764 100644
--- a/src/node/server/index.ts
+++ b/src/node/server/index.ts
@@ -36,6 +36,7 @@ export interface ServerPluginContext {
   watcher: HMRWatcher
   resolver: InternalResolver
   config: ServerConfig & { __path?: string }
+  port: number
 }
 
 export interface State extends DefaultState {}
@@ -71,7 +72,10 @@ export function createServer(config: ServerConfig): Server {
     server,
     watcher,
     resolver,
-    config
+    config,
+    // port is exposed on the context for hmr client connection
+    // in case the files are served under a different port
+    port: config.port || 3000
   }
 
   // attach server context to koa context
@@ -109,11 +113,12 @@ export function createServer(config: ServerConfig): Server {
   resolvedPlugins.forEach((m) => m && m(context))
 
   const listen = server.listen.bind(server)
-  server.listen = (async (...args: any[]) => {
+  server.listen = (async (port: number, ...args: any[]) => {
     if (optimizeDeps.auto !== false) {
       await require('../optimizer').optimizeDeps(config)
     }
-    return listen(...args)
+    context.port = port
+    return listen(port, ...args)
   }) as any
 
   return server
diff --git a/src/node/server/serverPluginHmr.ts b/src/node/server/serverPluginHmr.ts
index 8273b1abef2fed..f3c36db195a3cf 100644
--- a/src/node/server/serverPluginHmr.ts
+++ b/src/node/server/serverPluginHmr.ts
@@ -83,7 +83,7 @@ export const hmrPlugin: ServerPlugin = ({
     if (ctx.path === hmrClientPublicPath) {
       ctx.type = 'js'
       ctx.status = 200
-      ctx.body = hmrClient
+      ctx.body = hmrClient.replace(`__PORT__`, ctx.port.toString())
     } else {
       if (ctx.query.t) {
         latestVersionsMap.set(ctx.path, ctx.query.t)