From 56e93dcd5a3ae2d14fe28fc09bdaeba8e410dfd6 Mon Sep 17 00:00:00 2001 From: Hans Ott Date: Tue, 27 Feb 2024 16:53:24 +0100 Subject: [PATCH 1/4] Hide require-in-the-middle and Shimmer libraries This gives us a clean API to wrap things from a package: Allows for wrapping multiple packages at once and even based on versions --- library/src/agent/Wrapper.ts | 173 +++++++++++++++++++++++- library/src/agent/protect.ts | 40 ++---- library/src/sinks/MongoDB.test.ts | 5 +- library/src/sinks/MongoDB.ts | 198 +++++++++++++--------------- library/src/sinks/Postgres.test.ts | 5 +- library/src/sinks/Postgres.ts | 78 +++++------ library/src/sources/Express.test.ts | 5 +- library/src/sources/Express.ts | 41 ++---- 8 files changed, 324 insertions(+), 221 deletions(-) diff --git a/library/src/agent/Wrapper.ts b/library/src/agent/Wrapper.ts index 3fa1f6300..6055d5e21 100644 --- a/library/src/agent/Wrapper.ts +++ b/library/src/agent/Wrapper.ts @@ -1,3 +1,174 @@ +import { Hook } from "require-in-the-middle"; +import { wrap } from "shimmer"; +import { getPackageVersion } from "../helpers/getPackageVersion"; +import { satisfiesVersion } from "../helpers/satisfiesVersion"; + +type Interceptor = (args: unknown[], subject: unknown) => void | unknown[]; + +class Method { + constructor( + private readonly name: string, + private readonly interceptor: Interceptor + ) { + if (!this.name) { + throw new Error("Method name is required"); + } + } + + getName() { + return this.name; + } + + getInterceptor() { + return this.interceptor; + } +} + +class Selector { + private methods: Method[] = []; + + constructor(private readonly selector: (exports: unknown) => unknown) {} + + method(name: string, interceptor: Interceptor) { + const method = new Method(name, interceptor); + this.methods.push(method); + + return method; + } + + getSelector() { + return this.selector; + } + + getMethods() { + return this.methods; + } +} + +class VersionedPackage { + private selectors: Selector[] = []; + + constructor(private readonly range: string) { + if (!this.range) { + throw new Error("Version range is required"); + } + } + + getRange() { + return this.range; + } + + subject(selector: (exports: any) => unknown): Selector { + const fn = new Selector(selector); + this.selectors.push(fn); + + return fn; + } + + getSelectors() { + return this.selectors; + } +} + +class Package { + private versions: VersionedPackage[] = []; + + constructor(private readonly packageName: string) {} + + getName() { + return this.packageName; + } + + withVersion(range: string): VersionedPackage { + const pkg = new VersionedPackage(range); + this.versions.push(pkg); + + return pkg; + } + + getVersions() { + return this.versions; + } +} + +export class Hooks { + private readonly packages: Package[] = []; + + package(packageName: string): Package { + if (!packageName) { + throw new Error("Package name is required"); + } + + const pkg = new Package(packageName); + this.packages.push(pkg); + + return pkg; + } + + getPackages() { + return this.packages; + } +} + +export function wrapPackages(hooks: Hooks) { + const wrapped: Record = {}; + + hooks.getPackages().forEach((pkg) => { + const version = getPackageVersion(pkg.getName()); + + if (!version) { + return; + } + + const selectors = pkg + .getVersions() + .map((versioned) => { + if (!satisfiesVersion(versioned.getRange(), version)) { + return []; + } + + return versioned.getSelectors(); + }) + .flat(); + + if (selectors.length === 0) { + return; + } + + wrapped[pkg.getName()] = { + version, + supported: true, + }; + + new Hook([pkg.getName()], (exports) => { + selectors.forEach((selector) => { + const subject = selector.getSelector()(exports); + + if (!subject) { + return; + } + + selector.getMethods().forEach((method) => { + wrap(subject, method.getName(), function (original) { + return function (this: unknown, ...args: unknown[]) { + const updatedArgs = method.getInterceptor()(args, this); + + return original.apply( + this, + Array.isArray(updatedArgs) ? updatedArgs : args + ); + }; + }); + }); + }); + + return exports; + }); + }); + + return wrapped; +} + export interface Wrapper { - wrap(): void; + wrap(hooks: Hooks): void; } diff --git a/library/src/agent/protect.ts b/library/src/agent/protect.ts index e9b10d37a..05989c437 100644 --- a/library/src/agent/protect.ts +++ b/library/src/agent/protect.ts @@ -1,6 +1,4 @@ import type { APIGatewayProxyHandler } from "aws-lambda"; -import { getPackageVersion } from "../helpers/getPackageVersion"; -import { satisfiesVersion } from "../helpers/satisfiesVersion"; import { Agent } from "./Agent"; import { getInstance, setInstance } from "./AgentSingleton"; import { API, APIFetch, APIThrottled, Token } from "./API"; @@ -10,40 +8,18 @@ import { MongoDB } from "../sinks/MongoDB"; import { Postgres } from "../sinks/Postgres"; import * as shimmer from "shimmer"; import { Logger, LoggerConsole, LoggerNoop } from "./Logger"; -import { Wrapper } from "./Wrapper"; +import { Hooks, wrapPackages } from "./Wrapper"; import { Options, getOptions } from "../helpers/getOptions"; function wrapInstalledPackages() { - const packages: Record = { - express: { - range: "^4.0.0", - wrapper: new Express(), - }, - mongodb: { - range: "^4.0.0 || ^5.0.0 || ^6.0.0", - wrapper: new MongoDB(), - }, - pg: { - range: "^8.0.0", - wrapper: new Postgres(), - }, - }; - - const wrapped: Record = {}; - for (const packageName in packages) { - const { range, wrapper } = packages[packageName]; - const version = getPackageVersion(packageName); - wrapped[packageName] = { - version, - supported: version ? satisfiesVersion(range, version) : false, - }; - - if (wrapped[packageName].supported) { - wrapper.wrap(); - } - } + const wrappers = [new Express(), new MongoDB(), new Postgres()]; + + const hooks = new Hooks(); + wrappers.forEach((wrapper) => { + wrapper.wrap(hooks); + }); - return wrapped; + return wrapPackages(hooks); } function getLogger(options: Options): Logger { diff --git a/library/src/sinks/MongoDB.test.ts b/library/src/sinks/MongoDB.test.ts index f5d3fe739..465c4d611 100644 --- a/library/src/sinks/MongoDB.test.ts +++ b/library/src/sinks/MongoDB.test.ts @@ -4,6 +4,7 @@ import { setInstance } from "../agent/AgentSingleton"; import { APIForTesting, Token } from "../agent/API"; import { LoggerNoop } from "../agent/Logger"; import { Context, runWithContext } from "../agent/Context"; +import { Hooks, wrapPackages } from "../agent/Wrapper"; import { MongoDB } from "./MongoDB"; const context: Context = { @@ -21,7 +22,9 @@ const context: Context = { }; t.test("we can highjack the MongoDB library", async (t) => { - new MongoDB().wrap(); + const hooks = new Hooks(); + new MongoDB().wrap(hooks); + wrapPackages(hooks); const agent = new Agent( true, diff --git a/library/src/sinks/MongoDB.ts b/library/src/sinks/MongoDB.ts index d86fc3ea3..63f3ed50c 100644 --- a/library/src/sinks/MongoDB.ts +++ b/library/src/sinks/MongoDB.ts @@ -1,14 +1,12 @@ /* eslint-disable prefer-rest-params */ import { Collection } from "mongodb"; -import { Hook } from "require-in-the-middle"; -import { wrap } from "shimmer"; import { Agent } from "../agent/Agent"; import { getInstance } from "../agent/AgentSingleton"; import { detectNoSQLInjection } from "../vulnerabilities/nosql-injection/detectNoSQLInjection"; import { isPlainObject } from "../helpers/isPlainObject"; import { Context, getContext } from "../agent/Context"; import { friendlyName } from "../agent/Source"; -import { Wrapper } from "../agent/Wrapper"; +import { Hooks, Wrapper } from "../agent/Wrapper"; const OPERATIONS_WITH_FILTER = [ "count", @@ -82,121 +80,101 @@ export class MongoDB implements Wrapper { } } - private wrapBulkWrite(exports: unknown) { - const that = this; - - wrap( - // @ts-expect-error This is magic that TypeScript doesn't understand - exports.Collection.prototype, - "bulkWrite", - function wrapBulkWrite(original) { - return function protectBulkWrite(this: Collection) { - const agent = getInstance(); - - if (!agent) { - return original.apply(this, arguments); - } - - const request = getContext(); - - if (!request) { - agent.onInspectedCall({ - module: "mongodb", - withoutContext: true, - detectedAttack: false, - }); - - return original.apply(this, arguments); - } - - if (!Array.isArray(arguments[0])) { - return original.apply(this, arguments); - } - - const operations: BulkWriteOperation[] = arguments[0]; - operations.forEach((operation) => { - BULK_WRITE_OPERATIONS_WITH_FILTER.forEach((command) => { - const options = operation[command]; - - if (options && options.filter) { - that.inspectFilter( - this.dbName, - this.collectionName, - agent, - request, - options.filter, - "bulkWrite" - ); - } - }); - }); - - return original.apply(this, arguments); - }; - } - ); - } + private inspectBulkWrite(args: unknown[], collection: Collection) { + const agent = getInstance(); - private wrapOperationsWithFilter(exports: unknown) { - const that = this; + if (!agent) { + return; + } - OPERATIONS_WITH_FILTER.forEach((operation) => { - wrap( - // @ts-expect-error This is magic that TypeScript doesn't understand - exports.Collection.prototype, - operation, - function wrapOperation(original) { - return function protectQuery(this: Collection) { - const agent = getInstance(); - - if (!agent) { - return original.apply(this, arguments); - } - - const hasFilter = - arguments.length > 0 && isPlainObject(arguments[0]); - - if (!hasFilter) { - return original.apply(this, arguments); - } - - const request = getContext(); - - if (!request) { - agent.onInspectedCall({ - module: "mongodb", - withoutContext: true, - detectedAttack: false, - }); - - return original.apply(this, arguments); - } - - const filter = arguments[0]; - that.inspectFilter( - this.dbName, - this.collectionName, - agent, - request, - filter, - operation - ); - - return original.apply(this, arguments); - }; + const request = getContext(); + + if (!request) { + return agent.onInspectedCall({ + module: "mongodb", + withoutContext: true, + detectedAttack: false, + }); + } + + if (!Array.isArray(args[0])) { + return; + } + + const operations: BulkWriteOperation[] = args[0]; + operations.forEach((operation) => { + BULK_WRITE_OPERATIONS_WITH_FILTER.forEach((command) => { + const options = operation[command]; + + if (options && options.filter) { + this.inspectFilter( + collection.dbName, + collection.collectionName, + agent, + request, + options.filter, + "bulkWrite" + ); } - ); + }); }); } - private onModuleRequired(exports: T): T { - this.wrapBulkWrite(exports); - this.wrapOperationsWithFilter(exports); + private inspectOperation( + operation: string, + args: unknown[], + collection: Collection + ): void { + const agent = getInstance(); + + if (!agent) { + return; + } + + const hasFilter = args.length > 0 && isPlainObject(args[0]); - return exports; + if (!hasFilter) { + return; + } + + const request = getContext(); + + if (!request) { + return agent.onInspectedCall({ + module: "mongodb", + withoutContext: true, + detectedAttack: false, + }); + } + + const filter = args[0]; + this.inspectFilter( + collection.dbName, + collection.collectionName, + agent, + request, + filter, + operation + ); } - wrap() { - new Hook(["mongodb"], this.onModuleRequired.bind(this)); + wrap(hooks: Hooks) { + const mongodb = hooks + .package("mongodb") + .withVersion("^4.0.0 || ^5.0.0 || ^6.0.0"); + + const collection = mongodb.subject( + (exports) => exports.Collection.prototype + ); + + OPERATIONS_WITH_FILTER.forEach((operation) => { + collection.method(operation, (args, collection) => + this.inspectOperation(operation, args, collection as Collection) + ); + }); + + collection.method("bulkWrite", (args, collection) => + this.inspectBulkWrite(args, collection as Collection) + ); } } diff --git a/library/src/sinks/Postgres.test.ts b/library/src/sinks/Postgres.test.ts index 1660261a9..d3357d0cf 100644 --- a/library/src/sinks/Postgres.test.ts +++ b/library/src/sinks/Postgres.test.ts @@ -4,6 +4,7 @@ import { setInstance } from "../agent/AgentSingleton"; import { APIForTesting, Token } from "../agent/API"; import { LoggerNoop } from "../agent/Logger"; import { runWithContext, type Context } from "../agent/Context"; +import { Hooks, wrapPackages } from "../agent/Wrapper"; import { Postgres } from "./Postgres"; import type { Client } from "pg"; @@ -28,7 +29,9 @@ const context: Context = { }; t.test("We can hijack Postgres class", async () => { - new Postgres().wrap(); + const hooks = new Hooks(); + new Postgres().wrap(hooks); + wrapPackages(hooks); const agent = new Agent( true, diff --git a/library/src/sinks/Postgres.ts b/library/src/sinks/Postgres.ts index afde1e7fc..f1172b3f4 100644 --- a/library/src/sinks/Postgres.ts +++ b/library/src/sinks/Postgres.ts @@ -1,58 +1,42 @@ -import { Wrapper } from "../agent/Wrapper"; -import { Client } from "pg"; -import { Hook } from "require-in-the-middle"; -import { massWrap } from "shimmer"; +import { Hooks, Wrapper } from "../agent/Wrapper"; import { getInstance } from "../agent/AgentSingleton"; import { getContext } from "../agent/Context"; import { checkContextForSqlInjection } from "../vulnerabilities/sql-injection/detectSQLInjection"; export class Postgres implements Wrapper { - private wrapQueryFunction(exports: unknown) { - massWrap( - // @ts-expect-error This is magic that TypeScript doesn't understand - [exports.Client.prototype, exports.Pool.prototype], - ["query"], - function wrapQueryFunction(original) { - return function safeQueryFunction(this: Client, ...args: unknown[]) { - const agent = getInstance(); - - if (!agent) { - return original.apply(this, args); - } - - const request = getContext(); - - if (!request) { - agent.onInspectedCall({ - module: "postgres", - withoutContext: true, - detectedAttack: false, - }); - - return original.apply(this, args); - } - - if (typeof args[0] !== "string") { - // The query is not a string, not much to do here - return original.apply(this, args); - } - - const querystring: string = args[0]; - checkContextForSqlInjection(querystring, request, agent, "postgres"); - - return original.apply(this, args); - }; - } - ); + private inspectQuery(args: unknown[]) { + const agent = getInstance(); + + if (!agent) { + return; + } + + const request = getContext(); + + if (!request) { + return agent.onInspectedCall({ + module: "postgres", + withoutContext: true, + detectedAttack: false, + }); + } + + if (typeof args[0] !== "string") { + // The query is not a string, not much to do here + return; + } + + const querystring: string = args[0]; + checkContextForSqlInjection(querystring, request, agent, "postgres"); } - private onModuleRequired(exports: T): T { - this.wrapQueryFunction(exports); + wrap(hooks: Hooks) { + const pg = hooks.package("pg").withVersion("^7.0.0 || ^8.0.0"); - return exports; - } + const client = pg.subject((exports) => exports.Client.prototype); + client.method("query", (args) => this.inspectQuery(args)); - wrap() { - new Hook(["pg"], this.onModuleRequired.bind(this)); + const pool = pg.subject((exports) => exports.Pool.prototype); + pool.method("query", (args) => this.inspectQuery(args)); } } diff --git a/library/src/sources/Express.test.ts b/library/src/sources/Express.test.ts index 9449c68b9..6a6a944d2 100644 --- a/library/src/sources/Express.test.ts +++ b/library/src/sources/Express.test.ts @@ -1,8 +1,11 @@ import * as t from "tap"; +import { Hooks, wrapPackages } from "../agent/Wrapper"; import { Express } from "./Express"; // Before express is required! -new Express().wrap(); +const hooks = new Hooks(); +new Express().wrap(hooks); +wrapPackages(hooks); import * as express from "express"; import * as request from "supertest"; diff --git a/library/src/sources/Express.ts b/library/src/sources/Express.ts index c5ee73e6d..b2af7225b 100644 --- a/library/src/sources/Express.ts +++ b/library/src/sources/Express.ts @@ -1,9 +1,7 @@ /* eslint-disable prefer-rest-params */ -import type { NextFunction, Request, Response, Application } from "express"; -import { Hook } from "require-in-the-middle"; -import { massWrap } from "shimmer"; +import type { NextFunction, Request, Response } from "express"; import { runWithContext } from "../agent/Context"; -import { Wrapper } from "../agent/Wrapper"; +import { Hooks, Wrapper } from "../agent/Wrapper"; import { METHODS } from "node:http"; type Middleware = (req: Request, resp: Response, next: NextFunction) => void; @@ -40,33 +38,20 @@ export class Express implements Wrapper { // app.get("/path", json(), middleware(), (req, res) => { ... })) // // Without having to change the user's code - private wrapRouteMethods(exports: unknown) { - massWrap( - // @ts-expect-error This is magic that TypeScript doesn't understand - exports.Route.prototype, - // @ts-expect-error This is magic that TypeScript doesn't understand - METHODS.map((method) => method.toLowerCase()), - function wrapRouteMethod(original) { - return function injectMiddleware(this: Application) { - const args = Array.from(arguments); - const handler = args.pop(); - args.push(createMiddleware()); - args.push(handler); + private addMiddleware(args: unknown[]) { + const handler = args.pop(); + args.push(createMiddleware()); + args.push(handler); - // @ts-expect-error This is magic that TypeScript doesn't understand - return original.apply(this, args); - }; - } - ); + return args; } - private onModuleRequire(exports: T): T { - this.wrapRouteMethods(exports); - - return exports; - } + wrap(hooks: Hooks) { + const express = hooks.package("express").withVersion("^4.0.0"); - wrap() { - new Hook(["express"], this.onModuleRequire.bind(this)); + const route = express.subject((exports) => exports.Route.prototype); + METHODS.map((method) => method.toLowerCase()).forEach((method) => { + route.method(method, (args) => this.addMiddleware(args)); + }); } } From 890c4ec2477055adb00aa30d143a35ce2488a2c7 Mon Sep 17 00:00:00 2001 From: Hans Ott Date: Tue, 27 Feb 2024 17:09:55 +0100 Subject: [PATCH 2/4] Extract wrapMethod function --- library/src/agent/Wrapper.ts | 32 ++++++++++++++++++++------------ 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/library/src/agent/Wrapper.ts b/library/src/agent/Wrapper.ts index 6055d5e21..8a9350c98 100644 --- a/library/src/agent/Wrapper.ts +++ b/library/src/agent/Wrapper.ts @@ -110,6 +110,25 @@ export class Hooks { } } +function wrapMethod(subject: unknown, method: Method) { + // @ts-expect-error We don't now the type of the subject + wrap(subject, method.getName(), function wrap(original: Function) { + return function wrap() { + // eslint-disable-next-line prefer-rest-params + const args = Array.from(arguments); + // @ts-expect-error We don't now the type of this + const updatedArgs = method.getInterceptor()(args, this); + + return original.apply( + // @ts-expect-error We don't now the type of this + this, + // eslint-disable-next-line prefer-rest-params + Array.isArray(updatedArgs) ? updatedArgs : arguments + ); + }; + }); +} + export function wrapPackages(hooks: Hooks) { const wrapped: Record = {}; @@ -148,18 +167,7 @@ export function wrapPackages(hooks: Hooks) { return; } - selector.getMethods().forEach((method) => { - wrap(subject, method.getName(), function (original) { - return function (this: unknown, ...args: unknown[]) { - const updatedArgs = method.getInterceptor()(args, this); - - return original.apply( - this, - Array.isArray(updatedArgs) ? updatedArgs : args - ); - }; - }); - }); + selector.getMethods().forEach((method) => wrapMethod(subject, method)); }); return exports; From 6f0ad88ab8016c88ba682f881bf5e6d325d51590 Mon Sep 17 00:00:00 2001 From: Hans Ott Date: Tue, 27 Feb 2024 17:21:13 +0100 Subject: [PATCH 3/4] Bump code coverage --- library/src/agent/Wrapper.test.ts | 25 ++++++++++ library/src/agent/Wrapper.ts | 74 +---------------------------- library/src/agent/protect.ts | 5 +- library/src/agent/wrap.test.ts | 30 ++++++++++++ library/src/agent/wrap.ts | 72 ++++++++++++++++++++++++++++ library/src/sinks/MongoDB.test.ts | 5 +- library/src/sinks/Postgres.test.ts | 5 +- library/src/sources/Express.test.ts | 5 +- 8 files changed, 140 insertions(+), 81 deletions(-) create mode 100644 library/src/agent/Wrapper.test.ts create mode 100644 library/src/agent/wrap.test.ts create mode 100644 library/src/agent/wrap.ts diff --git a/library/src/agent/Wrapper.test.ts b/library/src/agent/Wrapper.test.ts new file mode 100644 index 000000000..ca08ae5f5 --- /dev/null +++ b/library/src/agent/Wrapper.test.ts @@ -0,0 +1,25 @@ +import * as t from "tap"; +import { Hooks } from "./Wrapper"; + +t.test("package throws error if name is empty", async (t) => { + const hooks = new Hooks(); + + t.throws(() => hooks.package("")); +}); + +t.test("withVersion throws error if version is empty", async (t) => { + const hooks = new Hooks(); + const subject = hooks.package("package"); + + t.throws(() => subject.withVersion("")); +}); + +t.test("method throws error if name is empty", async (t) => { + const hooks = new Hooks(); + const subject = hooks + .package("package") + .withVersion("^1.0.0") + .subject((exports) => exports); + + t.throws(() => subject.method("", () => {})); +}); diff --git a/library/src/agent/Wrapper.ts b/library/src/agent/Wrapper.ts index 8a9350c98..dde4e08eb 100644 --- a/library/src/agent/Wrapper.ts +++ b/library/src/agent/Wrapper.ts @@ -1,11 +1,6 @@ -import { Hook } from "require-in-the-middle"; -import { wrap } from "shimmer"; -import { getPackageVersion } from "../helpers/getPackageVersion"; -import { satisfiesVersion } from "../helpers/satisfiesVersion"; - type Interceptor = (args: unknown[], subject: unknown) => void | unknown[]; -class Method { +export class Method { constructor( private readonly name: string, private readonly interceptor: Interceptor @@ -110,73 +105,6 @@ export class Hooks { } } -function wrapMethod(subject: unknown, method: Method) { - // @ts-expect-error We don't now the type of the subject - wrap(subject, method.getName(), function wrap(original: Function) { - return function wrap() { - // eslint-disable-next-line prefer-rest-params - const args = Array.from(arguments); - // @ts-expect-error We don't now the type of this - const updatedArgs = method.getInterceptor()(args, this); - - return original.apply( - // @ts-expect-error We don't now the type of this - this, - // eslint-disable-next-line prefer-rest-params - Array.isArray(updatedArgs) ? updatedArgs : arguments - ); - }; - }); -} - -export function wrapPackages(hooks: Hooks) { - const wrapped: Record = {}; - - hooks.getPackages().forEach((pkg) => { - const version = getPackageVersion(pkg.getName()); - - if (!version) { - return; - } - - const selectors = pkg - .getVersions() - .map((versioned) => { - if (!satisfiesVersion(versioned.getRange(), version)) { - return []; - } - - return versioned.getSelectors(); - }) - .flat(); - - if (selectors.length === 0) { - return; - } - - wrapped[pkg.getName()] = { - version, - supported: true, - }; - - new Hook([pkg.getName()], (exports) => { - selectors.forEach((selector) => { - const subject = selector.getSelector()(exports); - - if (!subject) { - return; - } - - selector.getMethods().forEach((method) => wrapMethod(subject, method)); - }); - - return exports; - }); - }); - - return wrapped; -} - export interface Wrapper { wrap(hooks: Hooks): void; } diff --git a/library/src/agent/protect.ts b/library/src/agent/protect.ts index 05989c437..83d9cab27 100644 --- a/library/src/agent/protect.ts +++ b/library/src/agent/protect.ts @@ -8,7 +8,8 @@ import { MongoDB } from "../sinks/MongoDB"; import { Postgres } from "../sinks/Postgres"; import * as shimmer from "shimmer"; import { Logger, LoggerConsole, LoggerNoop } from "./Logger"; -import { Hooks, wrapPackages } from "./Wrapper"; +import { wrap } from "./wrap"; +import { Hooks } from "./Wrapper"; import { Options, getOptions } from "../helpers/getOptions"; function wrapInstalledPackages() { @@ -19,7 +20,7 @@ function wrapInstalledPackages() { wrapper.wrap(hooks); }); - return wrapPackages(hooks); + return wrap(hooks); } function getLogger(options: Options): Logger { diff --git a/library/src/agent/wrap.test.ts b/library/src/agent/wrap.test.ts new file mode 100644 index 000000000..15281f1cc --- /dev/null +++ b/library/src/agent/wrap.test.ts @@ -0,0 +1,30 @@ +import * as t from "tap"; +import { wrap } from "./wrap"; +import { Hooks } from "./Wrapper"; + +t.test("it ignores if package is not installed", async (t) => { + const hooks = new Hooks(); + hooks.package("unknown").withVersion("^1.0.0"); + t.same(wrap(hooks), {}); +}); + +t.test("it ignores if packages have empty selectors", async (t) => { + const hooks = new Hooks(); + hooks.package("shimmer").withVersion("^1.0.0"); + t.same(wrap(hooks), {}); +}); + +t.test("it ignores unknown selectors", async (t) => { + const hooks = new Hooks(); + hooks + .package("shimmer") + .withVersion("^1.0.0") + .subject((exports) => exports.doesNotExist) + .method("method", () => {}); + t.same(wrap(hooks), { + shimmer: { + version: "1.2.1", + supported: true, + }, + }); +}); diff --git a/library/src/agent/wrap.ts b/library/src/agent/wrap.ts new file mode 100644 index 000000000..21bac6342 --- /dev/null +++ b/library/src/agent/wrap.ts @@ -0,0 +1,72 @@ +import { Hook } from "require-in-the-middle"; +import { wrap as shim } from "shimmer"; +import { getPackageVersion } from "../helpers/getPackageVersion"; +import { satisfiesVersion } from "../helpers/satisfiesVersion"; +import { Hooks, Method } from "./Wrapper"; + +function wrapMethod(subject: unknown, method: Method) { + // @ts-expect-error We don't now the type of the subject + shim(subject, method.getName(), function wrap(original: Function) { + return function wrap() { + // eslint-disable-next-line prefer-rest-params + const args = Array.from(arguments); + // @ts-expect-error We don't now the type of this + const updatedArgs = method.getInterceptor()(args, this); + + return original.apply( + // @ts-expect-error We don't now the type of this + this, + // eslint-disable-next-line prefer-rest-params + Array.isArray(updatedArgs) ? updatedArgs : arguments + ); + }; + }); +} + +export function wrap(hooks: Hooks) { + const wrapped: Record = {}; + + hooks.getPackages().forEach((pkg) => { + const version = getPackageVersion(pkg.getName()); + + if (!version) { + return; + } + + const selectors = pkg + .getVersions() + .map((versioned) => { + if (!satisfiesVersion(versioned.getRange(), version)) { + return []; + } + + return versioned.getSelectors(); + }) + .flat(); + + if (selectors.length === 0) { + return; + } + + wrapped[pkg.getName()] = { + version, + supported: true, + }; + + new Hook([pkg.getName()], (exports) => { + selectors.forEach((selector) => { + const subject = selector.getSelector()(exports); + + if (!subject) { + return; + } + + selector.getMethods().forEach((method) => wrapMethod(subject, method)); + }); + + return exports; + }); + }); + + return wrapped; +} diff --git a/library/src/sinks/MongoDB.test.ts b/library/src/sinks/MongoDB.test.ts index 465c4d611..4bed039e4 100644 --- a/library/src/sinks/MongoDB.test.ts +++ b/library/src/sinks/MongoDB.test.ts @@ -4,7 +4,8 @@ import { setInstance } from "../agent/AgentSingleton"; import { APIForTesting, Token } from "../agent/API"; import { LoggerNoop } from "../agent/Logger"; import { Context, runWithContext } from "../agent/Context"; -import { Hooks, wrapPackages } from "../agent/Wrapper"; +import { wrap } from "../agent/wrap"; +import { Hooks } from "../agent/Wrapper"; import { MongoDB } from "./MongoDB"; const context: Context = { @@ -24,7 +25,7 @@ const context: Context = { t.test("we can highjack the MongoDB library", async (t) => { const hooks = new Hooks(); new MongoDB().wrap(hooks); - wrapPackages(hooks); + wrap(hooks); const agent = new Agent( true, diff --git a/library/src/sinks/Postgres.test.ts b/library/src/sinks/Postgres.test.ts index d3357d0cf..7480b4337 100644 --- a/library/src/sinks/Postgres.test.ts +++ b/library/src/sinks/Postgres.test.ts @@ -4,7 +4,8 @@ import { setInstance } from "../agent/AgentSingleton"; import { APIForTesting, Token } from "../agent/API"; import { LoggerNoop } from "../agent/Logger"; import { runWithContext, type Context } from "../agent/Context"; -import { Hooks, wrapPackages } from "../agent/Wrapper"; +import { wrap } from "../agent/wrap"; +import { Hooks } from "../agent/Wrapper"; import { Postgres } from "./Postgres"; import type { Client } from "pg"; @@ -31,7 +32,7 @@ const context: Context = { t.test("We can hijack Postgres class", async () => { const hooks = new Hooks(); new Postgres().wrap(hooks); - wrapPackages(hooks); + wrap(hooks); const agent = new Agent( true, diff --git a/library/src/sources/Express.test.ts b/library/src/sources/Express.test.ts index 6a6a944d2..25970ac56 100644 --- a/library/src/sources/Express.test.ts +++ b/library/src/sources/Express.test.ts @@ -1,11 +1,12 @@ import * as t from "tap"; -import { Hooks, wrapPackages } from "../agent/Wrapper"; +import { wrap } from "../agent/wrap"; +import { Hooks } from "../agent/Wrapper"; import { Express } from "./Express"; // Before express is required! const hooks = new Hooks(); new Express().wrap(hooks); -wrapPackages(hooks); +wrap(hooks); import * as express from "express"; import * as request from "supertest"; From f06c022b43327d25e3db677fabcc9b740aa010aa Mon Sep 17 00:00:00 2001 From: Hans Ott Date: Tue, 27 Feb 2024 17:27:32 +0100 Subject: [PATCH 4/4] Bump code coverage --- library/src/agent/wrap.test.ts | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/library/src/agent/wrap.test.ts b/library/src/agent/wrap.test.ts index 15281f1cc..d8553f540 100644 --- a/library/src/agent/wrap.test.ts +++ b/library/src/agent/wrap.test.ts @@ -5,12 +5,14 @@ import { Hooks } from "./Wrapper"; t.test("it ignores if package is not installed", async (t) => { const hooks = new Hooks(); hooks.package("unknown").withVersion("^1.0.0"); + t.same(wrap(hooks), {}); }); t.test("it ignores if packages have empty selectors", async (t) => { const hooks = new Hooks(); hooks.package("shimmer").withVersion("^1.0.0"); + t.same(wrap(hooks), {}); }); @@ -21,10 +23,25 @@ t.test("it ignores unknown selectors", async (t) => { .withVersion("^1.0.0") .subject((exports) => exports.doesNotExist) .method("method", () => {}); + t.same(wrap(hooks), { shimmer: { version: "1.2.1", supported: true, }, }); + + // Force require to load shimmer + require("shimmer"); +}); + +t.test("it ignores if version is not supported", async (t) => { + const hooks = new Hooks(); + hooks + .package("shimmer") + .withVersion("^2.0.0") + .subject((exports) => exports) + .method("method", () => {}); + + t.same(wrap(hooks), {}); });