Refactor mysql2 wrapping

library/sinks/MySQL2.tests.ts

@@ -34,7 +34,7 @@ export function createMySQL2Tests(versionPkgName: string) {
   };
 
   t.test("it detects SQL injections", async (t) => {
-    const agent = startTestAgent({
+    startTestAgent({
       wrappers: [new MySQL2()],
       rewrite: {
         mysql2: versionPkgName,
@@ -54,17 +54,9 @@ export function createMySQL2Tests(versionPkgName: string) {
       multipleStatements: true,
     });
 
-    const mysqlCallback = require(
-      versionPkgName
-    ) as typeof import("mysql2-v3.12");
-    const connection2 = mysqlCallback.createConnection({
-      host: "localhost",
-      user: "root",
-      password: "mypassword",
-      database: "catsdb",
-      port: 27015,
-      multipleStatements: true,
-    });
+    let connection2:
+      | ReturnType<typeof import("mysql2-v3.12").createConnection>
+      | undefined;
 
     try {
       await connection.query(
@@ -128,10 +120,24 @@ export function createMySQL2Tests(versionPkgName: string) {
         return connection.execute("SELECT 1");
       });
 
+      // !!! Do not move this code up
+      // Because the connection of mysql2/promises will also be wrapped and possible test failures if only /promise is imported will be hidden
+      const mysqlCallback = require(
+        versionPkgName
+      ) as typeof import("mysql2-v3.12");
+      connection2 = mysqlCallback.createConnection({
+        host: "localhost",
+        user: "root",
+        password: "mypassword",
+        database: "catsdb",
+        port: 27015,
+        multipleStatements: true,
+      });
+
       const error3 = await t.rejects(async () => {
         await runWithContext(dangerousContext, () => {
           return new Promise((resolve, reject) => {
-            connection2.query(
+            connection2!.query(
               "-- should be blocked",
               (error: any, results: any) => {
                 if (error) {
@@ -152,13 +158,15 @@ export function createMySQL2Tests(versionPkgName: string) {
       }
 
       runWithContext(safeContext, () => {
-        connection2.query("-- This is a comment");
+        connection2!.query("-- This is a comment");
       });
     } catch (error: any) {
       t.fail(error);
     } finally {
       await connection.end();
-      await connection2.end();
+      if (connection2) {
+        await connection2.end();
+      }
     }
   });
 }

library/sinks/MySQL2.ts

@@ -2,8 +2,10 @@ import { getContext } from "../agent/Context";
 import { Hooks } from "../agent/hooks/Hooks";
 import { InterceptorResult } from "../agent/hooks/InterceptorResult";
 import { wrapExport } from "../agent/hooks/wrapExport";
+import { WrapPackageInfo } from "../agent/hooks/WrapPackageInfo";
 import { Wrapper } from "../agent/Wrapper";
 import { isPlainObject } from "../helpers/isPlainObject";
+import { isWrapped } from "../helpers/wrap";
 import { checkContextForSqlInjection } from "../vulnerabilities/sql-injection/checkContextForSqlInjection";
 import { SQLDialect } from "../vulnerabilities/sql-injection/dialects/SQLDialect";
 import { SQLDialectMySQL } from "../vulnerabilities/sql-injection/dialects/SQLDialectMySQL";
@@ -49,36 +51,65 @@ export class MySQL2 implements Wrapper {
     return undefined;
   }
 
+  // This function is copied from the OpenTelemetry MySQL2 instrumentation (Apache 2.0 license)
+  // https://github.com/open-telemetry/opentelemetry-js-contrib/blob/21e1331a29e06092fb1e460ca99e0c28b1b57ac4/plugins/node/opentelemetry-instrumentation-mysql2/src/utils.ts#L150
+  private getConnectionPrototypeToInstrument(connection: any) {
+    const connectionPrototype = connection.prototype;
+    const basePrototype = Object.getPrototypeOf(connectionPrototype);
+
+    // [email protected] included a refactoring, where most code was moved out of the `Connection` class and into a shared base
+    // so we need to instrument that instead, see https://github.com/sidorares/node-mysql2/pull/3081
+    // This checks if the functions we're instrumenting are there on the base - we cannot use the presence of a base
+    // prototype since EventEmitter is the base for mysql2@<=3.11.4
+    if (
+      typeof basePrototype?.query === "function" &&
+      typeof basePrototype?.execute === "function"
+    ) {
+      return basePrototype;
+    }
+
+    // otherwise instrument the connection directly.
+    return connectionPrototype;
+  }
+
   wrap(hooks: Hooks) {
+    const wrapConnection = (
+      exports: any,
+      pkgInfo: WrapPackageInfo,
+      isPromise: boolean
+    ) => {
+      const connectionPrototype = this.getConnectionPrototypeToInstrument(
+        isPromise ? exports.PromiseConnection : exports.Connection
+      );
+
+      if (!isWrapped(connectionPrototype.query)) {
+        // Wrap connection.query
+        wrapExport(connectionPrototype, "query", pkgInfo, {
+          inspectArgs: (args, agent) => this.inspectQuery("mysql2.query", args),
+        });
+      }
+
+      if (!isWrapped(connectionPrototype.execute)) {
+        // Wrap connection.execute
+        wrapExport(connectionPrototype, "execute", pkgInfo, {
+          inspectArgs: (args, agent) =>
+            this.inspectQuery("mysql2.execute", args),
+        });
+      }
+    };
+
     const pkg = hooks.addPackage("mysql2");
     // For all versions of mysql2 newer than 3.0.0
-    pkg.withVersion("^3.0.0").onRequire((exports, pkgInfo) => {
-      // Wrap connection.query
-      wrapExport(exports.Connection.prototype, "query", pkgInfo, {
-        inspectArgs: (args, agent) => this.inspectQuery("mysql2.query", args),
-      });
-
-      // Wrap connection.execute
-      wrapExport(exports.Connection.prototype, "execute", pkgInfo, {
-        inspectArgs: (args, agent) => this.inspectQuery("mysql2.execute", args),
-      });
-    });
+    pkg
+      .withVersion("^3.0.0")
+      .onRequire((exports, pkgInfo) => wrapConnection(exports, pkgInfo, false));
 
     // For all versions of mysql2 newer than / equal 3.11.5
     // Reason: https://github.com/sidorares/node-mysql2/pull/3081
     pkg
       .withVersion("^3.11.5")
       .onFileRequire("promise.js", (exports, pkgInfo) => {
-        // Wrap PromiseConnection.query
-        wrapExport(exports.PromiseConnection.prototype, "query", pkgInfo, {
-          inspectArgs: (args, agent) => this.inspectQuery("mysql2.query", args),
-        });
-
-        // Wrap PromiseConnection.execute
-        wrapExport(exports.PromiseConnection.prototype, "execute", pkgInfo, {
-          inspectArgs: (args, agent) =>
-            this.inspectQuery("mysql2.execute", args),
-        });
+        return wrapConnection(exports, pkgInfo, true);
       });
   }
 }