From cfcf10c1f3bd3a33b253e647a997917458a0d5b3 Mon Sep 17 00:00:00 2001
From: Embbnux Ji <embbnux@embbnux.com>
Date: Sat, 28 May 2022 23:51:46 +0800
Subject: [PATCH] feat: add x-forwarded header (#5)

---
 server.js | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/server.js b/server.js
index 545d889..de6c68b 100644
--- a/server.js
+++ b/server.js
@@ -75,6 +75,24 @@ app.get('/tunnel_jwt_generator', (req, res) => {
   res.send('Forbidden');
 });
 
+function getReqHeaders(req) {
+  const encrypted = !!(req.isSpdy || req.connection.encrypted || req.connection.pair);
+  const headers = { ...req.headers };
+  const url = new URL(`${encrypted ? 'https' : 'http'}://${req.headers.host}`);
+  const forwardValues = {
+    for: req.connection.remoteAddress || req.socket.remoteAddress,
+    port: url.port || (encrypted ? 443 : 80),
+    proto: encrypted ? 'https' : 'http',
+  };
+  ['for', 'port', 'proto'].forEach((key) => {
+    const previousValue = req.headers[`x-forwarded-${key}`] || '';
+    headers[`x-forwarded-${key}`] =
+      `${previousValue || ''}${previousValue ? ',' : ''}${forwardValues[key]}`;
+  });
+  headers['x-forwarded-host'] = req.headers['x-forwarded-host'] || req.headers.host || '';
+  return headers;
+}
+
 app.use('/', (req, res) => {
   const tunnelSocket = tunnelSockets[req.headers.host];
   if (!tunnelSocket) {
@@ -88,7 +106,7 @@ app.use('/', (req, res) => {
     requestId,
     request: {
       method: req.method,
-      headers: { ...req.headers },
+      headers: getReqHeaders(req),
       path: req.url,
     },
   });
@@ -164,7 +182,7 @@ httpServer.on('upgrade', (req, socket, head) => {
     requestId,
     request: {
       method: req.method,
-      headers: { ...req.headers },
+      headers: getReqHeaders(req),
       path: req.url,
     },
   });