core: add wrapMethod/unwrapMethod util

Author: isaacs

packages/core/src/utils/object.ts

@@ -18,6 +18,7 @@ import { isElement, isError, isEvent, isInstanceOf, isPrimitive } from './is';
  * args>)` or `origMethod.apply(this, [<other args>])` (rather than being called directly), again to preserve `this`.
  * @returns void
  */
+
 export function fill(source: { [key: string]: any }, name: string, replacementFactory: (...args: any[]) => any): void {
   if (!(name in source)) {
     return;
@@ -80,6 +81,36 @@ export function markFunctionWrapped(wrapped: WrappedFunction, original: WrappedF
   } catch {} // eslint-disable-line no-empty
 }
 
+/**
+ * Wrap a method on an object by name, only if it is not already wrapped.
+ */
+export function wrapMethod<O extends {}, T extends string & keyof O>(obj: O, field: T, wrapped: WrappedFunction): void {
+  const original = obj[field];
+  if (typeof original !== 'function') {
+    throw new Error(`Cannot wrap method: ${field} is not a function`);
+  }
+  if (getOriginalFunction(original)) {
+    throw new Error(`Attempting to wrap method ${field} multiple times`);
+  }
+  markFunctionWrapped(wrapped, original);
+  addNonEnumerableProperty(obj, field, wrapped);
+}
+
+/**
+ * Unwrap a previously wrapped method on an object by name
+ */
+export function unwrapMethod<O extends {}, T extends string & keyof O>(obj: O, field: T): void {
+  const wrapped = obj[field];
+  if (typeof wrapped !== 'function') {
+    throw new Error(`Cannot unwrap method: ${field} is not a function`);
+  }
+  const original = getOriginalFunction(wrapped);
+  if (!original) {
+    throw new Error(`Method ${field} is not wrapped, and cannot be unwrapped`);
+  }
+  addNonEnumerableProperty(obj, field, original);
+}
+
 /**
  * This extracts the original function if available.  See
  * `markFunctionWrapped` for more information.

packages/core/test/lib/utils/object.test.ts

@@ -11,6 +11,8 @@ import {
   fill,
   markFunctionWrapped,
   objectify,
+  unwrapMethod,
+  wrapMethod,
 } from '../../../src/utils/object';
 import { testOnlyIfNodeVersionAtLeast } from '../../testutils';
 
@@ -455,3 +457,38 @@ describe('markFunctionWrapped', () => {
     expect(originalFunc).not.toHaveBeenCalled();
   });
 });
+
+describe('unwrapMethod, wrapMethod', () => {
+  it('can wrap a method on an object and unwrap it later', () => {
+    const wrapped = () => {};
+    const original = () => {};
+    const obj = { m: original };
+    wrapMethod(obj, 'm', wrapped);
+    expect(obj.m).toBe(wrapped);
+    expect((obj.m as WrappedFunction).__sentry_original__).toBe(original);
+    unwrapMethod(obj, 'm');
+    expect(obj.m).toBe(original);
+  });
+
+  it('throws if misused', () => {
+    const wrapped = () => {};
+    const original = () => {};
+    const obj = { m: original };
+    wrapMethod(obj, 'm', wrapped);
+    expect(() => {
+      //@ts-expect-error verify type checking prevents this mistake
+      wrapMethod(obj, 'foo', wrapped);
+    }).toThrowError('Cannot wrap method: foo is not a function');
+    expect(() => {
+      //@ts-expect-error verify type checking prevents this mistake
+      unwrapMethod(obj, 'foo');
+    }).toThrowError('Cannot unwrap method: foo is not a function');
+    expect(() => {
+      wrapMethod(obj, 'm', wrapped);
+    }).toThrowError('Attempting to wrap method m multiple times');
+    unwrapMethod(obj, 'm');
+    expect(() => {
+      unwrapMethod(obj, 'm');
+    }).toThrowError('Method m is not wrapped, and cannot be unwrapped');
+  });
+});