From 4d5bccdfcb1cd347f25ce3a62e4eb599ac6cab9d Mon Sep 17 00:00:00 2001
From: Ben Allfree <ben@benallfree.com>
Date: Tue, 23 Jul 2024 21:01:15 +0000
Subject: [PATCH] feat(pockethost): sass custom domains support

---
 cspell.json                                   |  2 +
 packages/pockethost/src/constants.ts          |  3 +
 .../pockethost/src/services/ProxyService.ts   | 56 ++++++++++++++++---
 3 files changed, 52 insertions(+), 9 deletions(-)

diff --git a/cspell.json b/cspell.json
index d019bd8f..29bb7237 100644
--- a/cspell.json
+++ b/cspell.json
@@ -16,10 +16,12 @@
     "esbuild",
     "eventsource",
     "fullchain",
+    "geoip",
     "getenv",
     "gobot",
     "goja",
     "IPCIDR",
+    "ipcountry",
     "jsonifiable",
     "Jsonifiable",
     "lemonbot",
diff --git a/packages/pockethost/src/constants.ts b/packages/pockethost/src/constants.ts
index fb8aab9a..0501490d 100644
--- a/packages/pockethost/src/constants.ts
+++ b/packages/pockethost/src/constants.ts
@@ -107,6 +107,7 @@ export const SETTINGS = {
   PH_FTP_PASV_PORT_MIN: mkNumber(10000),
   PH_FTP_PASV_PORT_MAX: mkNumber(20000),
 
+  EDGE_SASS_DOMAINS_AUTH_TOKEN: mkString(``),
   EDGE_APEX_DOMAIN: mkString(_APEX_DOMAIN),
   EDGE_MAX_ACTIVE_INSTANCES: mkNumber(20),
 
@@ -232,6 +233,8 @@ export const PH_FTP_PASV_IP = () => settings().PH_FTP_PASV_IP
 export const PH_FTP_PASV_PORT_MIN = () => settings().PH_FTP_PASV_PORT_MIN
 export const PH_FTP_PASV_PORT_MAX = () => settings().PH_FTP_PASV_PORT_MAX
 
+export const EDGE_SASS_DOMAINS_AUTH_TOKEN = () =>
+  settings().EDGE_SASS_DOMAINS_AUTH_TOKEN
 export const EDGE_APEX_DOMAIN = () => settings().EDGE_APEX_DOMAIN
 export const EDGE_MAX_ACTIVE_INSTANCES = () =>
   settings().EDGE_MAX_ACTIVE_INSTANCES
diff --git a/packages/pockethost/src/services/ProxyService.ts b/packages/pockethost/src/services/ProxyService.ts
index 1fe96087..b9c8ee5b 100644
--- a/packages/pockethost/src/services/ProxyService.ts
+++ b/packages/pockethost/src/services/ProxyService.ts
@@ -5,6 +5,7 @@ import { default as Server, default as httpProxy } from 'http-proxy'
 import { AsyncReturnType } from 'type-fest'
 import {
   DAEMON_PORT,
+  EDGE_SASS_DOMAINS_AUTH_TOKEN,
   Logger,
   LoggerService,
   SingletonBaseConfig,
@@ -49,22 +50,17 @@ export const proxyService = mkSingleton(async (config: ProxyServiceConfig) => {
     res.end
   })
 
+  // Default locals
   server.use((req, res, next) => {
     const host = req.headers.host
-    if (!host) {
-      throw new Error(`Host not found`)
-    }
     res.locals.host = host
-    next()
-  })
-
-  server.use((req, res, next) => {
     res.locals.coreInternalUrl = coreInternalUrl
     next()
   })
 
+  // Cloudflare signature
   server.use((req, res, next) => {
-    const url = new URL(`http://${req.headers.host}${req.url}`)
+    const url = new URL(`https://${res.locals.host}${req.url}`)
     const country = (req.headers['cf-ipcountry'] as string) || '<ct>'
     const ip = (req.headers['x-forwarded-for'] as string) || '<ip>'
     const method = req.method || '<m>'
@@ -74,7 +70,36 @@ export const proxyService = mkSingleton(async (config: ProxyServiceConfig) => {
       ip.padEnd(45),
       url.toString(),
     ].join(' ')
-    info(`Incoming request ${sig}`)
+    res.locals.sig = sig
+    next()
+  })
+
+  // SaaS domains overrides
+  server.use((req, res, next) => {
+    if (!(`x-saas-domains-auth-token` in req.headers)) {
+      next()
+      return
+    }
+
+    const secret = EDGE_SASS_DOMAINS_AUTH_TOKEN()
+    if (req.headers[`x-saas-domains-auth-token`] !== secret) {
+      throw new Error(`Invalid SaaS domain secret`)
+    }
+
+    const host = req.headers[`x-served-for`]
+    res.locals.host = host
+
+    const url = new URL(`https://${host}${req.url}`)
+    const country =
+      (req.headers['x-saas-geoip-country-code'] as string) || '<ct>'
+    const ip = (req.headers['x-saas-domains-ip'] as string) || '<ip>'
+    const method = req.method || '<m>'
+    const sig = [
+      method.padStart(10),
+      country.padStart(5),
+      ip.padEnd(45),
+      url.toString(),
+    ].join(' ')
     res.locals.sig = sig
     next()
   })
@@ -84,6 +109,19 @@ export const proxyService = mkSingleton(async (config: ProxyServiceConfig) => {
     next()
   })
 
+  // Request logging
+  server.use((req, res, next) => {
+    if (!res.locals.host) {
+      throw new Error(`Host not found`)
+    }
+    next()
+  })
+
+  server.use((req, res, next) => {
+    info(`Incoming request ${res.locals.sig}`, req.headers)
+    next()
+  })
+
   server.listen(DAEMON_PORT(), () => {
     info(`daemon listening on port ${DAEMON_PORT()}`)
   })