Fix proxy ws and other behavior

packages/standalone-proxy/index.ts

@@ -6,6 +6,9 @@ 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'
+import { appLoggerFromEnv } from './src/logging'
+import pino from 'pino'
 
 const __dirname = url.fileURLToPath(new URL('.', import.meta.url))
 
@@ -32,8 +35,9 @@ const envStore = inMemoryPreviewEnvStore({
   },
 })
 
-const app = createApp({ envStore, sshPublicKey })
-const sshLogger = app.log.child({ name: 'ssh_server' })
+const logger = pino(appLoggerFromEnv())
+const app = createApp({ sshPublicKey, isProxyRequest: isProxyRequest(BASE_URL), proxyHandlers: proxyHandlers({envStore, logger}), logger })
+const sshLogger = logger.child({ name: 'ssh_server' })
 
 const tunnelName = (clientId: string, remotePath: string) => {
   const serviceName = remotePath.replace(/^\//, '')

packages/standalone-proxy/src/app.ts

@@ -1,39 +1,40 @@
-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'
+import {Logger} from 'pino'
 
-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, logger }: { 
   sshPublicKey: string
+  isProxyRequest: (req: http.IncomingMessage) => boolean
+  logger: Logger
+  proxyHandlers: { wsHandler: (req: http.IncomingMessage, socket: internal.Duplex, head: Buffer) => void, handler: (req: http.IncomingMessage, res: http.ServerResponse) => void }
 }) =>
   Fastify({
-    logger: appLoggerFromEnv(),
-    rewriteUrl,
+    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 {
+          logger.warn('unexpected upgrade request %j', {url: req.url, host: req.headers['host']})
+          socket.end()
+        }
+      })
+      return server;
+    },
+    logger,
   })
 
     .register(fastifyRequestContext)
     .get('/healthz', { logLevel: 'warn' }, async () => 'OK')
     .get('/ssh-public-key', async () => sshPublicKey)
-    .register(proxyRoutes, { prefix: '/proxy/', envStore })
 
 
 

packages/standalone-proxy/src/logging.ts

@@ -1,6 +1,6 @@
-import { FastifyBaseLogger, FastifyLoggerOptions, PinoLoggerOptions } from 'fastify/types/logger'
+import { PinoLoggerOptions } from 'fastify/types/logger'
 
-const envToLogger: Record<string, (FastifyLoggerOptions & PinoLoggerOptions) | FastifyBaseLogger | false> = {
+const envToLogger: Record<string, PinoLoggerOptions> = {
   development: {
     level: process.env.DEBUG ? 'debug' : 'info',
     transport: {
@@ -13,8 +13,7 @@ const envToLogger: Record<string, (FastifyLoggerOptions & PinoLoggerOptions) | F
   },
   production: {
     level: process.env.DEBUG ? 'debug' : 'info',
-  },
-  test: false,
+  }
 }
 
 export const appLoggerFromEnv = () => envToLogger[process.env.NODE_ENV || 'development']

packages/standalone-proxy/src/proxy.ts

@@ -1,38 +1,61 @@
-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'
+import type { Logger } from 'pino'
 
-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, 
+  logger
+}: {
+  envStore: PreviewEnvStore
+  logger: Logger
+} ){
+  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) {
+      logger.warn('no env for %j', { targetHost, url })
+      logger.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,
+      logger.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 +64,32 @@ export const proxyRoutes: FastifyPluginAsync<{ envStore: PreviewEnvStore }> = as
           },
         },
         (err) => {
-          req.log.warn('error in proxy %j', err, { targetHost, url })
+          logger.warn('error in proxy %j', { error:err, targetHost: env.target,  url: req.url })
+        }
+      )
+    }, (err)=> logger.error('error forwarding traffic %j', {error: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) => {
+          logger.warn('error in ws proxy %j', { error:err, targetHost: env.target,  url: req.url })
         }
       )
 
-      return res
-    },
-  })
-}
+    }, (err)=> logger.error('error forwarding ws traffic %j', {error: err}))
+  }
+}