feat: non existent domain support (#17)

* chore: added support for async onError handlers in proxy

* feat: added support for header and response modification when dns resolution fails

Author: wrongsahil

src/components/proxy-middleware/helpers/proxy_ctx_helper.js

@@ -20,7 +20,7 @@ export const get_original_request_headers = (ctx) => {
 
 export const get_original_response_headers = (ctx) => {
   // TODO: This needs to be fetched from ctx.clientToProxy headers
-  return ctx.serverToProxyResponse.headers;
+  return ctx?.serverToProxyResponse?.headers || {};
 };
 
 export const is_request_preflight = (ctx) => {

src/components/proxy-middleware/index.js

@@ -132,6 +132,39 @@ class ProxyMiddlewareManager {
         this.rulesHelper
       );
 
+      ctx.onError(async function (ctx, err, kind, callback) {
+        // Should only modify response body & headers
+        ctx.rq_response_body = "" + kind + ": " + err, "utf8";
+        const { action_result_objs, continue_request } = await rules_middleware.on_response(ctx);
+
+        // Only modify response if any modify_response action is applied
+        const modifyResponseActionExist = action_result_objs.some((action_result_obj) => action_result_obj?.action?.action === "modify_response")
+
+        if(modifyResponseActionExist) {
+          const statusCode = ctx.rq_response_status_code || 404;
+          const responseHeaders = getResponseHeaders(ctx) || {}
+          ctx.proxyToClientResponse.writeHead(
+              statusCode,
+              http.STATUS_CODES[statusCode],
+              responseHeaders
+            );
+          ctx.proxyToClientResponse.end(ctx.rq_response_body);
+
+          ctx.rq.set_final_response({
+              status_code: statusCode,
+              headers: responseHeaders,
+              body: ctx.rq_response_body,
+            });
+          logger_middleware.send_network_log(
+              ctx,
+              rules_middleware.action_result_objs,
+              GLOBAL_CONSTANTS.REQUEST_STATE.COMPLETE
+            )
+        }
+
+        return callback();
+      })
+
       let request_body_chunks = [];
       ctx.onRequestData(async function (ctx, chunk, callback) {
         if (chunk) {

src/components/proxy-middleware/rule_action_processor/processors/modify_header_processor.js

@@ -6,6 +6,7 @@ const process_modify_header_action = (action, ctx) => {
   const allowed_handlers = [
     PROXY_HANDLER_TYPE.ON_REQUEST,
     PROXY_HANDLER_TYPE.ON_RESPONSE,
+    PROXY_HANDLER_TYPE.ON_ERROR,
   ];
 
   if (!allowed_handlers.includes(ctx.currentHandler)) {
@@ -14,7 +15,7 @@ const process_modify_header_action = (action, ctx) => {
 
   if (ctx.currentHandler == PROXY_HANDLER_TYPE.ON_REQUEST) {
     modify_request_headers(action, ctx);
-  } else if (ctx.currentHandler == PROXY_HANDLER_TYPE.ON_RESPONSE) {
+  } else if (ctx.currentHandler === PROXY_HANDLER_TYPE.ON_RESPONSE || ctx.currentHandler === PROXY_HANDLER_TYPE.ON_ERROR) {
     modify_response_headers(action, ctx);
   }
   return build_action_processor_response(action, true);
@@ -35,6 +36,8 @@ const modify_request_headers = (action, ctx) => {
 };
 
 const modify_response_headers = (action, ctx) => {
+  ctx.serverToProxyResponse = ctx.serverToProxyResponse || {}
+  ctx.serverToProxyResponse.headers = ctx.serverToProxyResponse.headers || {}
   // {"header1":"val1", "header2":"val2"}
   const originalResponseHeadersObject = ctx.serverToProxyResponse.headers;
   //  ["header1","header2"]

src/components/proxy-middleware/rule_action_processor/processors/modify_response_processor.js

@@ -13,7 +13,7 @@ import { getFunctionFromString } from "../../../../utils";
 const { types } = require("util");
 
 const process_modify_response_action = async (action, ctx) => {
-  const allowed_handlers = [PROXY_HANDLER_TYPE.ON_RESPONSE_END];
+  const allowed_handlers = [PROXY_HANDLER_TYPE.ON_RESPONSE_END, PROXY_HANDLER_TYPE.ON_ERROR];
 
   if (!allowed_handlers.includes(ctx.currentHandler)) {
     return build_action_processor_response(action, false);
@@ -82,9 +82,9 @@ const modify_response_using_code = async (action, ctx) => {
           ? ctx.clientToProxyRequest.method
           : null
         : null,
-      response: ctx.rq_response_body,
+      response: ctx?.rq_response_body,
       url: get_request_url(ctx),
-      responseType: ctx.serverToProxyResponse.headers["content-type"],
+      responseType: ctx?.serverToProxyResponse?.headers?.["content-type"],
       requestHeaders: ctx.clientToProxyRequest.headers,
       requestData: parseJsonBody(ctx.rq?.final_request?.body) || null,
     };

src/lib/proxy/lib/proxy.ts

@@ -55,6 +55,7 @@ const PROXY_HANDLER_TYPE = {
   ON_RESPONSE_HEADERS: "ON_RESPONSE_HEADERS",
   ON_RESPONSE_DATA: "ON_RESPONSE_DATA",
   ON_RESPONSE_END: "ON_RESPONSE_END",
+  ON_ERROR: "ON_ERROR",
 };
 
 module.exports.PROXY_HANDLER_TYPE = PROXY_HANDLER_TYPE;
@@ -635,21 +636,29 @@ Proxy.prototype.onCertificateMissing = function (ctx, files, callback) {
 };
 
 Proxy.prototype._onError = function (kind, ctx, err) {
-  this.onErrorHandlers.forEach(function (handler) {
-    return handler(ctx, err, kind);
-  });
-  if (ctx) {
-    ctx.onErrorHandlers.forEach(function (handler) {
-      return handler(ctx, err, kind);
-    });
+  if(ctx) {
+    ctx.currentHandler = PROXY_HANDLER_TYPE.ON_ERROR;
+  }
 
-    if (ctx.proxyToClientResponse && !ctx.proxyToClientResponse.headersSent) {
-      ctx.proxyToClientResponse.writeHead(504, "Proxy Error");
-    }
-    if (ctx.proxyToClientResponse && !ctx.proxyToClientResponse.finished) {
-      ctx.proxyToClientResponse.end("" + kind + ": " + err, "utf8");
+  async.forEach(
+    this.onErrorHandlers.concat(ctx?.onErrorHandlers || []),
+    function (fn, callback) {
+      if (fn) {
+        return fn(ctx, err, kind, callback);
+      }
+      callback();
+    },
+    function() {
+      if(ctx) {
+        if (ctx.proxyToClientResponse && !ctx.proxyToClientResponse.headersSent) {
+          ctx.proxyToClientResponse.writeHead(504, "Proxy Error");
+        }
+        if (ctx.proxyToClientResponse && !ctx.proxyToClientResponse.finished) {
+          ctx.proxyToClientResponse.end("" + kind + ": " + err, "utf8");
+        }
+      }
     }
-  }
+  );
 };
 
 Proxy.prototype._onWebSocketServerConnect = function (isSSL, ws, upgradeReq) {