Fix bot and IP blocking for HTTP2 server in stream mode

library/sources/HTTP2Server.test.ts

@@ -8,6 +8,7 @@ import { isLocalhostIP } from "../helpers/isLocalhostIP";
 import { resolve } from "path";
 import { FileSystem } from "../sinks/FileSystem";
 import { createTestAgent } from "../helpers/createTestAgent";
+import { FetchListsAPIForTesting } from "../agent/api/FetchListsAPIForTesting";
 
 // Allow self-signed certificates
 process.env.NODE_TLS_REJECT_UNAUTHORIZED = "0";
@@ -42,9 +43,27 @@ const api = new ReportingAPIForTesting({
   ],
   heartbeatIntervalInMS: 10 * 60 * 1000,
 });
+
+const mockedFetchListAPI = new FetchListsAPIForTesting({
+  allowedIPAddresses: [],
+  blockedIPAddresses: [
+    {
+      key: "geoip/Belgium;BE",
+      source: "geoip",
+      ips: ["9.9.9.9"],
+      description: "geo restrictions",
+    },
+  ],
+  blockedUserAgents: "hackerbot",
+  monitoredUserAgents: "",
+  monitoredIPAddresses: [],
+  userAgentDetails: [],
+});
+
 const agent = createTestAgent({
   token: new Token("123"),
   api,
+  fetchListsAPI: mockedFetchListAPI,
 });
 agent.start([new HTTPServer(), new FileSystem()]);
 
@@ -821,3 +840,91 @@ t.test("invalid Multipart body results in empty body", async () => {
     });
   });
 });
+
+t.test("it blocks bots", async (t) => {
+  const server = http2.createSecureServer({
+    key: readFileSync(resolve(__dirname, "fixtures/key.pem")),
+    cert: readFileSync(resolve(__dirname, "fixtures/cert.pem")),
+  });
+
+  server.on("stream", (stream, headers) => {
+    stream.respond({ ":status": 200 });
+    stream.end(JSON.stringify(getContext()));
+  });
+
+  await new Promise<void>((resolve) => {
+    server.listen(3438, () => {
+      http2Request(new URL("https://localhost:3438"), "GET", {
+        "User-Agent": "hackerbot 1.0",
+      }).then(({ headers, body }) => {
+        t.same(headers[":status"], 403);
+        t.same(
+          body,
+          "You are not allowed to access this resource because you have been identified as a bot."
+        );
+
+        server.close();
+        resolve();
+      });
+    });
+  });
+});
+
+t.test("it blocks blocked IPs", async (t) => {
+  const server = http2.createSecureServer({
+    key: readFileSync(resolve(__dirname, "fixtures/key.pem")),
+    cert: readFileSync(resolve(__dirname, "fixtures/cert.pem")),
+  });
+
+  server.on("stream", (stream, headers) => {
+    stream.respond({ ":status": 200 });
+    stream.end(JSON.stringify(getContext()));
+  });
+
+  await new Promise<void>((resolve) => {
+    server.listen(3439, () => {
+      http2Request(new URL("https://localhost:3439"), "GET", {
+        "x-forwarded-for": "9.9.9.9",
+      }).then(({ headers, body }) => {
+        t.same(headers[":status"], 403);
+        t.same(
+          body,
+          "Your IP address is blocked due to geo restrictions. (Your IP: 9.9.9.9)"
+        );
+
+        server.close();
+        resolve();
+      });
+    });
+  });
+});
+
+t.test("it blocks blocked IPs using session stream event", async (t) => {
+  const server = http2.createSecureServer({
+    key: readFileSync(resolve(__dirname, "fixtures/key.pem")),
+    cert: readFileSync(resolve(__dirname, "fixtures/cert.pem")),
+  });
+
+  server.on("session", (session) => {
+    session.on("stream", (stream, headers) => {
+      stream.respond({ ":status": 200 });
+      stream.end(JSON.stringify(getContext()));
+    });
+  });
+
+  await new Promise<void>((resolve) => {
+    server.listen(3440, () => {
+      http2Request(new URL("https://localhost:3440"), "GET", {
+        "x-forwarded-for": "9.9.9.9",
+      }).then(({ headers, body }) => {
+        t.same(headers[":status"], 403);
+        t.same(
+          body,
+          "Your IP address is blocked due to geo restrictions. (Your IP: 9.9.9.9)"
+        );
+        server.close();
+        resolve();
+      });
+    });
+  });
+});

library/sources/HTTPServer.ts

@@ -3,6 +3,7 @@ import { Hooks } from "../agent/hooks/Hooks";
 import { wrapExport } from "../agent/hooks/wrapExport";
 import { Wrapper } from "../agent/Wrapper";
 import { createRequestListener } from "./http-server/createRequestListener";
+import { createSessionListener } from "./http-server/http2/createSessionListener";
 import { createStreamListener } from "./http-server/http2/createStreamListener";
 
 export class HTTPServer implements Wrapper {
@@ -39,6 +40,10 @@ export class HTTPServer implements Wrapper {
       return [args[0], createStreamListener(args[1], module, agent)];
     }
 
+    if (module === "http2" && args[0] === "session") {
+      return [args[0], createSessionListener(args[1], agent)];
+    }
+
     return args;
   }
 

library/sources/http-server/blockIPsAndBots.ts

@@ -3,6 +3,7 @@ import { Agent } from "../../agent/Agent";
 import { getContext } from "../../agent/Context";
 import { escapeHTML } from "../../helpers/escapeHTML";
 import { ipAllowedToAccessRoute } from "./ipAllowedToAccessRoute";
+import type { ServerHttp2Stream } from "http2";
 
 const checkedBlocks = Symbol("__zen_checked_blocks__");
 
@@ -16,7 +17,7 @@ export function blockIPsAndBots(
   // This flag is used to determine whether the request has already been checked
   // We use a Symbol so that we don't accidentally overwrite any other properties on the response object
   // and that we're the only ones that can access it
-  res: ServerResponse & { [checkedBlocks]?: boolean },
+  res: (ServerResponse | ServerHttp2Stream) & { [checkedBlocks]?: boolean },
   agent: Agent
 ): boolean {
   if (res.headersSent) {
@@ -40,15 +41,12 @@ export function blockIPsAndBots(
   res[checkedBlocks] = true;
 
   if (!ipAllowedToAccessRoute(context, agent)) {
-    res.statusCode = 403;
-    res.setHeader("Content-Type", "text/plain");
-
     let message = "Your IP address is not allowed to access this resource.";
     if (context.remoteAddress) {
       message += ` (Your IP: ${escapeHTML(context.remoteAddress)})`;
     }
 
-    res.end(message);
+    sendResponse(res, 403, message);
 
     return true;
   }
@@ -65,15 +63,12 @@ export function blockIPsAndBots(
     context.remoteAddress &&
     !agent.getConfig().isAllowedIPAddress(context.remoteAddress).allowed
   ) {
-    res.statusCode = 403;
-    res.setHeader("Content-Type", "text/plain");
-
     let message = "Your IP address is not allowed to access this resource.";
     if (context.remoteAddress) {
       message += ` (Your IP: ${escapeHTML(context.remoteAddress)})`;
     }
 
-    res.end(message);
+    sendResponse(res, 403, message);
 
     return true;
   }
@@ -99,15 +94,12 @@ export function blockIPsAndBots(
   }
 
   if (result.blocked) {
-    res.statusCode = 403;
-    res.setHeader("Content-Type", "text/plain");
-
     let message = `Your IP address is blocked due to ${escapeHTML(result.reason)}.`;
     if (context.remoteAddress) {
       message += ` (Your IP: ${escapeHTML(context.remoteAddress)})`;
     }
 
-    res.end(message);
+    sendResponse(res, 403, message);
 
     return true;
   }
@@ -136,10 +128,9 @@ export function blockIPsAndBots(
   }
 
   if (isUserAgentBlocked.blocked) {
-    res.statusCode = 403;
-    res.setHeader("Content-Type", "text/plain");
-
-    res.end(
+    sendResponse(
+      res,
+      403,
       "You are not allowed to access this resource because you have been identified as a bot."
     );
 
@@ -148,3 +139,24 @@ export function blockIPsAndBots(
 
   return false;
 }
+
+function isStream(
+  res: ServerResponse | ServerHttp2Stream
+): res is ServerHttp2Stream {
+  return "respond" in res;
+}
+
+function sendResponse(
+  res: ServerResponse | ServerHttp2Stream,
+  statusCode: number,
+  message: string
+) {
+  if (isStream(res)) {
+    res.respond({ ":status": statusCode, "Content-Type": "text/plain" });
+    res.end(message);
+    return;
+  }
+  res.statusCode = statusCode;
+  res.setHeader("Content-Type", "text/plain");
+  res.end(message);
+}

library/sources/http-server/http2/createSessionListener.ts

@@ -0,0 +1,44 @@
+import { Agent } from "../../../agent/Agent";
+import { ServerHttp2Session } from "http2";
+import { createStreamListener } from "./createStreamListener";
+
+/**
+ * Wraps the http2 session listener to be able to instrument stream events.
+ */
+export function createSessionListener(listener: Function, agent: Agent) {
+  return function sessionListener(session: ServerHttp2Session) {
+    // Wrap all session events to instrument stream events
+    session.on = wrapStreamEvent(session.on, agent);
+    session.once = wrapStreamEvent(session.once, agent);
+    session.addListener = wrapStreamEvent(session.addListener, agent);
+    session.prependListener = wrapStreamEvent(session.prependListener, agent);
+    session.prependOnceListener = wrapStreamEvent(
+      session.prependOnceListener,
+      agent
+    );
+
+    return listener(session);
+  };
+}
+
+function wrapStreamEvent(orig: Function, agent: Agent) {
+  return function wrap(...args: unknown[]) {
+    if (
+      args.length !== 2 ||
+      args[0] !== "stream" ||
+      typeof args[1] !== "function"
+    ) {
+      return orig.apply(
+        // @ts-expect-error We don't know the type of `this`
+        this,
+        arguments
+      );
+    }
+
+    return orig.apply(
+      // @ts-expect-error We don't know the type of `this`
+      this,
+      [args[0], createStreamListener(args[1], "http2", agent)]
+    );
+  };
+}

library/sources/http-server/http2/createStreamListener.ts

@@ -4,10 +4,12 @@ import {
   Context,
   getContext,
   runWithContext,
+  updateContext,
 } from "../../../agent/Context";
 import { contextFromStream } from "./contextFromStream";
 import { shouldDiscoverRoute } from "../shouldDiscoverRoute";
 import { IncomingHttpHeaders, ServerHttp2Stream } from "http2";
+import { blockIPsAndBots } from "../blockIPsAndBots";
 
 /**
  * Wraps the http2 stream listener to get the request context of http2 requests.
@@ -18,7 +20,7 @@ export function createStreamListener(
   module: string,
   agent: Agent
 ) {
-  return function requestListener(
+  return function streamListener(
     stream: ServerHttp2Stream,
     headers: IncomingHttpHeaders,
     flags: number,
@@ -43,6 +45,16 @@ export function createStreamListener(
       stream.prependListener = wrapStreamEvent(stream.prependListener);
       stream.prependOnceListener = wrapStreamEvent(stream.prependOnceListener);
 
+      if (blockIPsAndBots(stream, agent)) {
+        if (context) {
+          // To prevent attack wave detection from checking this request
+          updateContext(context, "blockedDueToIPOrBot", true);
+        }
+
+        // The return is necessary to prevent the listener from being called
+        return;
+      }
+
       return listener(stream, headers, flags, rawHeaders);
     });
   };