feat: upgrade bodyTransform middleware

Author: gkoos

README.md

@@ -78,6 +78,9 @@ global:
   - failRandomly:
       rate: 0.1
       status: 503
+  - bodyTransform:
+      request: "(body, ctx) => { body.foo = 'bar'; return body; }"
+      response: "(body, ctx) => { body.transformed = true; return body; }"
 routes:
   "GET /users/:id":
     - failRandomly:
@@ -120,7 +123,7 @@ Chaos Proxy uses Koa Router for path matching, supporting named parameters (e.g.
 - `rateLimit({ limit, windowMs, key })` — rate limiting (by IP, header, or custom)
 - `cors({ origin, methods, headers })` — enable and configure CORS headers. All options are strings.
 `throttle({ rate, chunkSize, burst, key })` — throttles bandwidth per request to a specified rate (bytes per second), with optional burst capacity and chunk size. The key option allows per-client throttling. (Implemented natively, not using koa-throttle.)
-- `bodyTransform({ transform })` — parse and mutate request body with a custom function.
+- `bodyTransform({ request?, response? })` — parse and mutate request and/or response body with custom functions.
 
 ### Rate Limiting
 
@@ -191,28 +194,28 @@ This configuration throttles responses to an average of 1 KB/s, sending data in
 
 ### Body Transform
 
-The `bodyTransform` middleware parses the incoming request body (JSON, form, or text), then replaces it with the result of your custom `transform` function. This lets you inspect, modify, or completely replace the body before it is proxied or processed by other middlewares. It uses `koa-bodyparser` under the hood.
+The `bodyTransform` middleware allows you to parse and mutate both the request and response bodies using custom transformation functions. You can specify a `request` and/or `response` transform, each as either a JavaScript function string (for YAML config) or a real function (for programmatic usage). Backward compatibility with the old `transform` key is removed—use the new object shape only.
 
 How it works:
 - Parses the request body and makes it available as `ctx.request.body`.
-- Calls your `transform` function with the parsed body and Koa context.
-- Sets the return value as the new `ctx.request.body`.
-- Subsequent middlewares and proxy logic use the mutated body.
+- If a `request` transform is provided, it is called with the parsed body and Koa context, and its return value replaces `ctx.request.body`.
+- After downstream middleware and proxying, if a `response` transform is provided, it is called with the response body (`ctx.body`) and Koa context, and its return value replaces `ctx.body`.
+- Both transforms can be used independently or together.
 
-**Example:**
+**Example (YAML):**
 ```yaml
 global:
   - bodyTransform:
-      transform: "(body, ctx) => { body.foo = 'bar'; return body; }"
+      request: "(body, ctx) => { body.foo = 'bar'; return body; }"
+      response: "(body, ctx) => { body.transformed = true; return body; }"
 ```
 
-This configuration adds a `foo: 'bar'` property to every JSON request body before it is proxied to the target server.
+This configuration adds a `foo: 'bar'` property to every JSON request body and a `transformed: true` property to every JSON response body.
 
 **Note:**
-For maximum flexibility, the `transform` option in `bodyTransform` can be specified as a JavaScript function string in your YAML config. This allows you to define custom transformation logic directly in the config file. Be aware that evaluating JS from config can introduce security and syntax risks. Use with care and only in trusted environments.
+For maximum flexibility, the `request` and `response` options in `bodyTransform` can be specified as JavaScript function strings in your YAML config. This allows you to define custom transformation logic directly in the config file. Be aware that evaluating JS from config can introduce security and syntax risks. Use with care and only in trusted environments.
 
-
-If you call `startServer` programmatically, you can also pass a real function instead of a string:
+If you call `startServer` programmatically, you can also pass real functions instead of strings:
 
 ```ts
 import { startServer, bodyTransform } from 'chaos-proxy';
@@ -222,9 +225,13 @@ startServer({
   port: 5000,
   global: [
     bodyTransform({
-      transform: (body, ctx) => {
+      request: (body, ctx) => {
         body.foo = 'bar';
         return body;
+      },
+      response: (body, ctx) => {
+        body.transformed = true;
+        return body;
       }
     })
   ]

src/middlewares/bodyTransform.ts

@@ -2,32 +2,64 @@ import bodyParser from 'koa-bodyparser';
 import type { Middleware, Context } from 'koa';
 
 export interface BodyTransformOptions {
-  transform: (body: unknown, ctx: Context) => unknown;
+  request?: { transform: (body: unknown, ctx: Context) => unknown } | string;
+  response?: { transform: (body: unknown, ctx: Context) => unknown } | string;
 }
 
-export function bodyTransform(opts: BodyTransformOptions | string): Middleware {
-  let transformFn: (body: unknown, ctx: Context) => unknown;
-  if (typeof opts === 'string') {
+function isTransformObject(opt: unknown): opt is { transform: (body: unknown, ctx: Context) => unknown } {
+  return (
+    typeof opt === 'object' &&
+    opt !== null &&
+    'transform' in opt &&
+    typeof (opt as { transform: unknown }).transform === 'function'
+  );
+}
+
+function parseTransform(opt: { transform: (body: unknown, ctx: Context) => unknown } | string | undefined, which: string): ((body: unknown, ctx: Context) => unknown) | undefined {
+  if (!opt) return undefined;
+  if (typeof opt === 'string') {
     try {
-      if (opts.trim().startsWith('(')) {
-        transformFn = eval(opts);
+      if (opt.trim().startsWith('(')) {
+        // Function string, e.g. (body, ctx) => ...
+        return eval(opt);
       } else {
-        transformFn = new Function('body', 'ctx', opts) as (body: unknown, ctx: Context) => unknown;
+        // Function body string, e.g. 'return {...body, foo: "bar"}'
+        return new Function('body', 'ctx', opt) as (body: unknown, ctx: Context) => unknown;
       }
     } catch (e) {
-      throw new Error('Failed to evaluate bodyTransform function string: ' + (e as Error).message);
+      throw new Error(`Failed to evaluate bodyTransform ${which} function string: ${(e as Error).message}`);
     }
+  } else if (isTransformObject(opt)) {
+    return opt.transform;
   } else {
-    transformFn = opts.transform;
+    throw new Error(`Invalid bodyTransform ${which} option: must be a function or string`);
+  }
+}
+
+export function bodyTransform(opts: BodyTransformOptions): Middleware {
+  if (typeof opts !== 'object' || (!opts.request && !opts.response)) {
+    throw new Error('bodyTransform expects an object with request and/or response keys');
   }
 
+  const requestTransform = parseTransform(opts.request, 'request');
+  const responseTransform = parseTransform(opts.response, 'response');
   const parser = bodyParser();
+
   return async (ctx: Context, next: () => Promise<void>) => {
-    await parser(ctx, async () => {
-      if (typeof transformFn === 'function' && ctx.request.body !== undefined) {
-        ctx.request.body = transformFn(ctx.request.body, ctx);
-      }
+    // Transform request body if needed
+    if (requestTransform) {
+      await parser(ctx, async () => {
+        if (ctx.request.body !== undefined) {
+          ctx.request.body = requestTransform(ctx.request.body, ctx);
+        }
+        await next();
+      });
+    } else {
       await next();
-    });
+    }
+    // Transform response body if needed
+    if (responseTransform && ctx.body !== undefined) {
+      ctx.body = responseTransform(ctx.body, ctx);
+    }
   };
 }

test/middlewares/bodyTransform.test.ts

@@ -16,53 +16,146 @@ function createMockCtx(body: unknown, contentType = 'application/json'): Context
 }
 
 describe('bodyTransform middleware', () => {
-  it('mutates JSON body', async () => {
+  it('mutates JSON request body', async () => {
     const mw = bodyTransform({
-      transform: (body) => {
-        if (typeof body === 'object' && body !== null) {
-          (body as Record<string, unknown>).mutated = true;
-        }
-        return body;
+      request: {
+        transform: (body: unknown) => {
+          if (typeof body === 'object' && body !== null) {
+            (body as Record<string, unknown>).mutated = true;
+          }
+          return body;
+        },
       },
     });
     const ctx = createMockCtx({ foo: 'bar' });
     await mw(ctx, async () => {});
     expect(ctx.request.body).toEqual({ foo: 'bar', mutated: true });
   });
 
-  it('handles non-object bodies', async () => {
+  it('handles non-object request bodies', async () => {
     const mw = bodyTransform({
-      transform: () => 'changed',
+      request: {
+  transform: () => 'changed',
+      },
     });
     const ctx = createMockCtx('original', 'text/plain');
     await mw(ctx, async () => {});
     expect(ctx.request.body).toBe('changed');
   });
 
-  it('returns undefined if transform returns undefined', async () => {
+  it('returns undefined if request transform returns undefined', async () => {
     const mw = bodyTransform({
-      transform: () => undefined,
+      request: {
+  transform: () => undefined,
+      },
     });
     const ctx = createMockCtx({ foo: 'bar' });
     await mw(ctx, async () => {});
     expect(ctx.request.body).toBeUndefined();
   });
 
-  it('accepts a string arrow function for transform', async () => {
-    const mw = bodyTransform('(body, ctx) => { body.added = 123; return body; }');
+  it('mutates response body', async () => {
+    const mw = bodyTransform({
+      response: {
+        transform: (body: unknown) => {
+          if (typeof body === 'object' && body !== null) {
+            (body as Record<string, unknown>).mutated = true;
+          }
+          return body;
+        },
+      },
+    });
+    const ctx = createMockCtx(undefined);
+    ctx.body = { foo: 'bar' };
+    await mw(ctx, async () => {});
+    expect(ctx.body).toEqual({ foo: 'bar', mutated: true });
+  });
+
+  it('handles non-object response bodies', async () => {
+    const mw = bodyTransform({
+      response: {
+  transform: () => 'changed',
+      },
+    });
+    const ctx = createMockCtx(undefined);
+    ctx.body = 'original';
+    await mw(ctx, async () => {});
+    expect(ctx.body).toBe('changed');
+  });
+
+  it('returns undefined if response transform returns undefined', async () => {
+    const mw = bodyTransform({
+      response: {
+  transform: () => undefined,
+      },
+    });
+    const ctx = createMockCtx(undefined);
+    ctx.body = { foo: 'bar' };
+    await mw(ctx, async () => {});
+    expect(ctx.body).toBeUndefined();
+  });
+
+  it('can use both request and response transforms', async () => {
+    const mw = bodyTransform({
+      request: {
+        transform: (body: unknown) => {
+          if (typeof body === 'object' && body !== null) {
+            (body as Record<string, unknown>).req = true;
+          }
+          return body;
+        },
+      },
+      response: {
+        transform: (body: unknown) => {
+          if (typeof body === 'object' && body !== null) {
+            (body as Record<string, unknown>).res = true;
+          }
+          return body;
+        },
+      },
+    });
+    const ctx = createMockCtx({ foo: 'bar' });
+    ctx.body = { bar: 'baz' };
+    await mw(ctx, async () => {});
+    expect(ctx.request.body).toEqual({ foo: 'bar', req: true });
+    expect(ctx.body).toEqual({ bar: 'baz', res: true });
+  });
+
+  it('accepts a string arrow function for request transform', async () => {
+    const mw = bodyTransform({ request: '(body, ctx) => { body.added = 123; return body; }' });
     const ctx = createMockCtx({ test: true });
     await mw(ctx, async () => {});
     expect(ctx.request.body).toEqual({ test: true, added: 123 });
   });
 
-  it('accepts a string function body for transform', async () => {
-    const mw = bodyTransform('body.foo = "baz"; return body;');
+  it('accepts a string function body for request transform', async () => {
+    const mw = bodyTransform({ request: 'body.foo = "baz"; return body;' });
     const ctx = createMockCtx({ foo: 'bar' });
     await mw(ctx, async () => {});
     expect(ctx.request.body).toEqual({ foo: 'baz' });
   });
 
-  it('throws for invalid function string', async () => {
-    expect(() => bodyTransform('not valid js')).toThrow();
+  it('accepts a string arrow function for response transform', async () => {
+    const mw = bodyTransform({ response: '(body, ctx) => { body.added = 456; return body; }' });
+    const ctx = createMockCtx(undefined);
+    ctx.body = { test: true };
+    await mw(ctx, async () => {});
+    expect(ctx.body).toEqual({ test: true, added: 456 });
+  });
+
+  it('accepts a string function body for response transform', async () => {
+    const mw = bodyTransform({ response: 'body.foo = "baz"; return body;' });
+    const ctx = createMockCtx(undefined);
+    ctx.body = { foo: 'bar' };
+    await mw(ctx, async () => {});
+    expect(ctx.body).toEqual({ foo: 'baz' });
+  });
+
+  it('throws for invalid function string in request', async () => {
+    expect(() => bodyTransform({ request: 'not valid js' })).toThrow();
+  });
+
+  it('throws for invalid function string in response', async () => {
+    expect(() => bodyTransform({ response: 'not valid js' })).toThrow();
   });
 });