smartproxy/ts/smartproxy.classes.networkproxy.ts

import * as plugins from './smartproxy.plugins.js';
import { ProxyRouter } from './smartproxy.classes.router.js';
import * as fs from 'fs';
import * as path from 'path';
import { fileURLToPath } from 'url';

export interface INetworkProxyOptions {
  port: number;
}

interface IWebSocketWithHeartbeat extends plugins.wsDefault {
  lastPong: number;
}

export class NetworkProxy {
  public options: INetworkProxyOptions;
  public proxyConfigs: plugins.tsclass.network.IReverseProxyConfig[] = [];
  public httpsServer: plugins.https.Server;
  public router = new ProxyRouter();
  public socketMap = new plugins.lik.ObjectMap<plugins.net.Socket>();
  public defaultHeaders: { [key: string]: string } = {};
  public heartbeatInterval: NodeJS.Timeout;
  private defaultCertificates: { key: string; cert: string };

  public alreadyAddedReverseConfigs: {
    [hostName: string]: plugins.tsclass.network.IReverseProxyConfig;
  } = {};

  constructor(optionsArg: INetworkProxyOptions) {
    this.options = optionsArg;
    const __dirname = path.dirname(fileURLToPath(import.meta.url));
    const certPath = path.join(__dirname, '..', 'assets', 'certs');
    
    try {
      this.defaultCertificates = {
        key: fs.readFileSync(path.join(certPath, 'key.pem'), 'utf8'),
        cert: fs.readFileSync(path.join(certPath, 'cert.pem'), 'utf8')
      };
    } catch (error) {
      console.error('Error loading certificates:', error);
      throw error;
    }
  }

  public async start() {
    // Instead of marking the callback async (which Node won't await),
    // we call our async handler and catch errors.
    this.httpsServer = plugins.https.createServer(
      {
        key: this.defaultCertificates.key,
        cert: this.defaultCertificates.cert
      },
      (originRequest, originResponse) => {
        this.handleRequest(originRequest, originResponse).catch((error) => {
          console.error('Unhandled error in request handler:', error);
          try {
            originResponse.end();
          } catch (err) {
            // ignore errors during cleanup
          }
        });
      },
    );

    // Enable websockets
    const wsServer = new plugins.ws.WebSocketServer({ server: this.httpsServer });

    // Set up the heartbeat interval
    this.heartbeatInterval = setInterval(() => {
      wsServer.clients.forEach((ws: plugins.wsDefault) => {
        const wsIncoming = ws as IWebSocketWithHeartbeat;
        if (!wsIncoming.lastPong) {
          wsIncoming.lastPong = Date.now();
        }
        if (Date.now() - wsIncoming.lastPong > 5 * 60 * 1000) {
          console.log('Terminating websocket due to missing pong for 5 minutes.');
          wsIncoming.terminate();
        } else {
          wsIncoming.ping();
        }
      });
    }, 60000); // runs every 1 minute

    wsServer.on(
      'connection',
      (wsIncoming: IWebSocketWithHeartbeat, reqArg: plugins.http.IncomingMessage) => {
        console.log(
          `wss proxy: got connection for wsc for https://${reqArg.headers.host}${reqArg.url}`,
        );

        wsIncoming.lastPong = Date.now();
        wsIncoming.on('pong', () => {
          wsIncoming.lastPong = Date.now();
        });

        let wsOutgoing: plugins.wsDefault;
        const outGoingDeferred = plugins.smartpromise.defer();

        // --- Improvement 2: Only call routeReq once ---
        const wsDestinationConfig = this.router.routeReq(reqArg);
        if (!wsDestinationConfig) {
          wsIncoming.terminate();
          return;
        }
        try {
          wsOutgoing = new plugins.wsDefault(
            `ws://${wsDestinationConfig.destinationIp}:${wsDestinationConfig.destinationPort}${reqArg.url}`,
          );
          console.log('wss proxy: initiated outgoing proxy');
          wsOutgoing.on('open', async () => {
            outGoingDeferred.resolve();
          });
        } catch (err) {
          console.error('Error initiating outgoing WebSocket:', err);
          wsIncoming.terminate();
          return;
        }

        wsIncoming.on('message', async (message, isBinary) => {
          try {
            await outGoingDeferred.promise;
            wsOutgoing.send(message, { binary: isBinary });
          } catch (error) {
            console.error('Error sending message to wsOutgoing:', error);
          }
        });

        wsOutgoing.on('message', async (message, isBinary) => {
          try {
            wsIncoming.send(message, { binary: isBinary });
          } catch (error) {
            console.error('Error sending message to wsIncoming:', error);
          }
        });

        const terminateWsOutgoing = () => {
          if (wsOutgoing) {
            wsOutgoing.terminate();
            console.log('Terminated outgoing ws.');
          }
        };
        wsIncoming.on('error', terminateWsOutgoing);
        wsIncoming.on('close', terminateWsOutgoing);

        const terminateWsIncoming = () => {
          if (wsIncoming) {
            wsIncoming.terminate();
            console.log('Terminated incoming ws.');
          }
        };
        wsOutgoing.on('error', terminateWsIncoming);
        wsOutgoing.on('close', terminateWsIncoming);
      },
    );

    this.httpsServer.keepAliveTimeout = 600 * 1000;
    this.httpsServer.headersTimeout = 600 * 1000;

    this.httpsServer.on('connection', (connection: plugins.net.Socket) => {
      this.socketMap.add(connection);
      console.log(`Added connection. Now ${this.socketMap.getArray().length} sockets connected.`);
      const cleanupConnection = () => {
        if (this.socketMap.checkForObject(connection)) {
          this.socketMap.remove(connection);
          console.log(`Removed connection. ${this.socketMap.getArray().length} sockets remaining.`);
          connection.destroy();
        }
      };
      connection.on('close', cleanupConnection);
      connection.on('error', cleanupConnection);
      connection.on('end', cleanupConnection);
      connection.on('timeout', cleanupConnection);
    });

    this.httpsServer.listen(this.options.port);
    console.log(
      `NetworkProxy -> OK: now listening for new connections on port ${this.options.port}`,
    );
  }

  /**
   * Internal async handler for processing HTTP/HTTPS requests.
   */
  private async handleRequest(
    originRequest: plugins.http.IncomingMessage,
    originResponse: plugins.http.ServerResponse,
  ): Promise<void> {
    const endOriginReqRes = (
      statusArg: number = 404,
      messageArg: string = 'This route is not available on this server.',
      headers: plugins.http.OutgoingHttpHeaders = {},
    ) => {
      originResponse.writeHead(statusArg, messageArg);
      originResponse.end(messageArg);
      if (originRequest.socket !== originResponse.socket) {
        console.log('hey, something is strange.');
      }
      originResponse.destroy();
    };

    console.log(
      `got request: ${originRequest.headers.host}${plugins.url.parse(originRequest.url).path}`,
    );
    const destinationConfig = this.router.routeReq(originRequest);

    if (!destinationConfig) {
      console.log(
        `${originRequest.headers.host} can't be routed properly. Terminating request.`,
      );
      endOriginReqRes();
      return;
    }

    // authentication
    if (destinationConfig.authentication) {
      const authInfo = destinationConfig.authentication;
      switch (authInfo.type) {
        case 'Basic': {
          const authHeader = originRequest.headers.authorization;
          if (!authHeader) {
            return endOriginReqRes(401, 'Authentication required', {
              'WWW-Authenticate': 'Basic realm="Access to the staging site", charset="UTF-8"',
            });
          }
          if (!authHeader.includes('Basic ')) {
            return endOriginReqRes(401, 'Authentication required', {
              'WWW-Authenticate': 'Basic realm="Access to the staging site", charset="UTF-8"',
            });
          }
          const authStringBase64 = authHeader.replace('Basic ', '');
          const authString: string = plugins.smartstring.base64.decode(authStringBase64);
          const userPassArray = authString.split(':');
          const user = userPassArray[0];
          const pass = userPassArray[1];
          if (user === authInfo.user && pass === authInfo.pass) {
            console.log('Request successfully authenticated');
          } else {
            return endOriginReqRes(403, 'Forbidden: Wrong credentials');
          }
          break;
        }
        default:
          return endOriginReqRes(
            403,
            'Forbidden: unsupported authentication method configured. Please report to the admin.',
          );
      }
    }

    let destinationUrl: string;
    if (destinationConfig) {
      destinationUrl = `http://${destinationConfig.destinationIp}:${destinationConfig.destinationPort}${originRequest.url}`;
    } else {
      return endOriginReqRes();
    }
    console.log(destinationUrl);
    try {
      const proxyResponse = await plugins.smartrequest.request(
        destinationUrl,
        {
          method: originRequest.method,
          headers: {
            ...originRequest.headers,
            'X-Forwarded-Host': originRequest.headers.host,
            'X-Forwarded-Proto': 'https',
          },
          keepAlive: true,
        },
        true, // streaming (keepAlive)
        (proxyRequest) => {
          originRequest.on('data', (data) => {
            proxyRequest.write(data);
          });
          originRequest.on('end', () => {
            proxyRequest.end();
          });
          originRequest.on('error', () => {
            proxyRequest.end();
          });
          originRequest.on('close', () => {
            proxyRequest.end();
          });
          originRequest.on('timeout', () => {
            proxyRequest.end();
            originRequest.destroy();
          });
          proxyRequest.on('error', () => {
            endOriginReqRes();
          });
        },
      );
      originResponse.statusCode = proxyResponse.statusCode;
      console.log(proxyResponse.statusCode);
      for (const defaultHeader of Object.keys(this.defaultHeaders)) {
        originResponse.setHeader(defaultHeader, this.defaultHeaders[defaultHeader]);
      }
      for (const header of Object.keys(proxyResponse.headers)) {
        originResponse.setHeader(header, proxyResponse.headers[header]);
      }
      proxyResponse.on('data', (data) => {
        originResponse.write(data);
      });
      proxyResponse.on('end', () => {
        originResponse.end();
      });
      proxyResponse.on('error', () => {
        originResponse.destroy();
      });
      proxyResponse.on('close', () => {
        originResponse.end();
      });
      proxyResponse.on('timeout', () => {
        originResponse.end();
        originResponse.destroy();
      });
    } catch (error) {
      console.error('Error while processing request:', error);
      endOriginReqRes(502, 'Bad Gateway: Error processing the request');
    }
  }

  public async updateProxyConfigs(
    proxyConfigsArg: plugins.tsclass.network.IReverseProxyConfig[],
  ) {
    console.log(`got new proxy configs`);
    this.proxyConfigs = proxyConfigsArg;
    this.router.setNewProxyConfigs(proxyConfigsArg);
    for (const hostCandidate of this.proxyConfigs) {
      const existingHostNameConfig = this.alreadyAddedReverseConfigs[hostCandidate.hostName];

      if (!existingHostNameConfig) {
        this.alreadyAddedReverseConfigs[hostCandidate.hostName] = hostCandidate;
      } else {
        if (
          existingHostNameConfig.publicKey === hostCandidate.publicKey &&
          existingHostNameConfig.privateKey === hostCandidate.privateKey
        ) {
          continue;
        } else {
          this.alreadyAddedReverseConfigs[hostCandidate.hostName] = hostCandidate;
        }
      }

      this.httpsServer.addContext(hostCandidate.hostName, {
        cert: hostCandidate.publicKey,
        key: hostCandidate.privateKey,
      });
    }
  }

  public async addDefaultHeaders(headersArg: { [key: string]: string }) {
    for (const headerKey of Object.keys(headersArg)) {
      this.defaultHeaders[headerKey] = headersArg[headerKey];
    }
  }

  public async stop() {
    const done = plugins.smartpromise.defer();
    this.httpsServer.close(() => {
      done.resolve();
    });
    for (const socket of this.socketMap.getArray()) {
      socket.destroy();
    }
    await done.promise;
    clearInterval(this.heartbeatInterval);
    console.log('NetworkProxy -> OK: Server has been stopped and all connections closed.');
  }
}