Fix proxy ws and other behavior

packages/standalone-proxy/index.ts

@@ -6,6 +6,7 @@ import { sshServer as createSshServer } from './src/ssh-server'
 import { getSSHKeys } from './src/ssh-keys'
 import url from 'url'
 import path from 'path'
+import { isProxyRequest, proxyHandlers } from './src/proxy'
 
 const __dirname = url.fileURLToPath(new URL('.', import.meta.url))
 
@@ -32,7 +33,8 @@ const envStore = inMemoryPreviewEnvStore({
   },
 })
 
-const app = createApp({ envStore, sshPublicKey })
+
+const app = createApp({ sshPublicKey, isProxyRequest: isProxyRequest(BASE_URL), proxyHandlers: proxyHandlers(envStore) })
 const sshLogger = app.log.child({ name: 'ssh_server' })
 
 const tunnelName = (clientId: string, remotePath: string) => {

packages/standalone-proxy/src/app.ts

@@ -1,39 +1,38 @@
-import Fastify, { RawRequestDefaultExpression } from 'fastify'
+import Fastify from 'fastify'
 import { fastifyRequestContext } from '@fastify/request-context'
-import { InternalServerError } from './errors'
 import { appLoggerFromEnv } from './logging'
-import { proxyRoutes } from './proxy'
-import { PreviewEnvStore } from './preview-env'
+import http from 'http'
+import internal from 'stream'
 
-const rewriteUrl = ({ url, headers: { host } }: RawRequestDefaultExpression): string => {
-  if (!url) {
-    throw new InternalServerError('no url in request')
-  }
-  if (!host) {
-    throw new InternalServerError('no host header in request')
-  }
-
-  const target = host.split('.', 1)[0]
-  if (!target.includes('-')) {
-    return url
-  }
-
-  return `/proxy/${target}${url}`
-}
-
-export const app = ({ envStore, sshPublicKey }: { 
-  envStore: PreviewEnvStore
+export const app = ({ sshPublicKey,isProxyRequest, proxyHandlers }: { 
   sshPublicKey: string
+  isProxyRequest: (req: http.IncomingMessage) => boolean
+  proxyHandlers: { wsHandler: (req: http.IncomingMessage, socket: internal.Duplex, head: Buffer) => void, handler: (req: http.IncomingMessage, res: http.ServerResponse) => void }
 }) =>
   Fastify({
+    serverFactory: (handler) => {
+      const {wsHandler:proxyWsHandler, handler: proxyHandler } = proxyHandlers
+      const server = http.createServer((req, res) => {
+        if (isProxyRequest(req)){
+          return proxyHandler(req, res)
+        }
+        return handler(req, res)
+      })
+      server.on('upgrade', (req, socket, head) => {
+        if (isProxyRequest(req)){
+          proxyWsHandler(req, socket, head)
+        } else {
+          socket.end()
+        }
+      })
+      return server;
+    },
     logger: appLoggerFromEnv(),
-    rewriteUrl,
   })
 
     .register(fastifyRequestContext)
     .get('/healthz', { logLevel: 'warn' }, async () => 'OK')
     .get('/ssh-public-key', async () => sshPublicKey)
-    .register(proxyRoutes, { prefix: '/proxy/', envStore })
 
 
 

packages/standalone-proxy/src/proxy.ts

@@ -1,38 +1,54 @@
-import { FastifyPluginAsync, HTTPMethods } from 'fastify'
 import { PreviewEnvStore } from './preview-env'
-import { NotFoundError } from './errors'
 import httpProxy from 'http-proxy'
+import { IncomingMessage, ServerResponse } from 'http'
+import internal from 'stream'
 
-const ALL_METHODS = Object.freeze(['DELETE', 'GET', 'HEAD', 'PATCH', 'POST', 'PUT', 'OPTIONS']) as HTTPMethods[]
-
-export const proxyRoutes: FastifyPluginAsync<{ envStore: PreviewEnvStore }> = async (app, { envStore }) => {
-  const proxy = httpProxy.createProxy({})
-
-  app.addHook('onClose', () => proxy.close())
-
-  // prevent FST_ERR_CTP_INVALID_MEDIA_TYPE error
-  app.removeAllContentTypeParsers()
-  app.addContentTypeParser('*', function (_request, _payload, done) { done(null) })
+export const isProxyRequest = (baseUrl: {hostname:string, port:string}) => (req: IncomingMessage)=> {
+  const host = req.headers["host"]
+  if (!host) return false
+  const {hostname: reqHostname, port: reqPort} = new URL(`http://${host}`)
+  if (reqPort !== baseUrl.port) return false
+  return reqHostname.endsWith(baseUrl.hostname) && reqHostname !== baseUrl.hostname
+}
 
-  app.route<{
-    Params: { targetHost: string; ['*']: string }
-  }>({
-    url: ':targetHost/*',
-    method: ALL_METHODS,
-    handler: async (req, res) => {
-      const { targetHost, ['*']: url } = req.params
-      req.log.debug('proxy request: %j', { targetHost, url, params: req.params })
-      const env = await envStore.get(targetHost)
+function asyncHandler<TArgs extends unknown[]>(fn: (...args: TArgs) => Promise<void>, onError: (error: unknown, ...args: TArgs)=> void ) {
+  return async (...args: TArgs) => {
+    try {
+      await fn(...args)
+    } catch (err) {
+      onError(err, ...args)
+    }
+  }
+}
 
+export function proxyHandlers(envStore: PreviewEnvStore, log=console){
+  const proxy = httpProxy.createProxy({})
+  const resolveTargetEnv = async (req: IncomingMessage)=>{
+    const {url} = req
+    const host = req.headers['host']
+    const targetHost = host?.split('.', 1)[0]
+    const env = await envStore.get(targetHost as string)
+    if (!env) {
+      log.warn('no env for %j', { targetHost, url })
+      log.warn('no host header in request')
+      return;
+    }
+    return env
+  }
+  return {
+    handler: asyncHandler(async (req: IncomingMessage, res: ServerResponse<IncomingMessage>) => {
+      const env = await resolveTargetEnv(req)
       if (!env) {
-        throw new NotFoundError(`host ${targetHost}`)
+        res.statusCode = 502;
+        res.end();
+        return;
       }
 
-      req.raw.url = `/${url}`
-
-      proxy.web(
-        req.raw,
-        res.raw,
+      log.info('proxying to %j', { target: env.target, url: req.url })
+      
+      return proxy.web(
+        req,
+        res,
         {
           // eslint-disable-next-line @typescript-eslint/ban-ts-comment
           // @ts-ignore
@@ -41,11 +57,32 @@ export const proxyRoutes: FastifyPluginAsync<{ envStore: PreviewEnvStore }> = as
           },
         },
         (err) => {
-          req.log.warn('error in proxy %j', err, { targetHost, url })
+          log.warn('error in proxy %j', err, { targetHost: env.target,  url: req.url })
+        }
+      )
+    }, (err)=> log.error('error in proxy %j', err) ),
+    wsHandler: asyncHandler(async (req: IncomingMessage, socket: internal.Duplex, head: Buffer) => {
+      const env = await resolveTargetEnv(req)
+      if (!env) {
+        socket.end();
+        return;
+      }
+      return proxy.ws(
+        req,
+        socket,
+        head,
+        {
+          // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+          // @ts-ignore
+          target: {
+            socketPath: env.target,
+          },
+        },
+        (err) => {
+          log.warn('error in proxy %j', err, { targetHost: env.target, url: req.url })
         }
       )
 
-      return res
-    },
-  })
-}
+    }, (err)=> log.error('error forwarding ws traffic %j', err))
+  }
+}