Fix bug with remote client use of updateHeader transforms

Author: pimterry

src/rules/requests/request-handlers.ts

@@ -656,12 +656,18 @@ interface SerializedPassThroughData {
     lookupOptions?: PassThroughLookupOptions;
 
     transformRequest?: Replace<
-        Replace<RequestTransform, 'replaceBody', string | undefined>, // Serialized as base64 buffer
-        'updateJsonBody', string | undefined // Serialized as a string to preserve undefined values
+        RequestTransform,
+        | 'replaceBody' // Serialized as base64 buffer
+        | 'updateHeaders' // // Serialized as a string to preserve undefined values
+        | 'updateJsonBody', // Serialized as a string to preserve undefined values
+        string | undefined
     >,
     transformResponse?: Replace<
-        Replace<ResponseTransform, 'replaceBody', string | undefined>, // Serialized as base64 buffer
-        'updateJsonBody', string | undefined // Serialized as a string to preserve undefined values
+        ResponseTransform,
+        | 'replaceBody' // Serialized as base64 buffer
+        | 'updateHeaders' // // Serialized as a string to preserve undefined values
+        | 'updateJsonBody', // Serialized as a string to preserve undefined values
+        string | undefined
     >,
 
     hasBeforeRequestCallback?: boolean;
@@ -852,6 +858,17 @@ function validateCustomHeaders(
 const OMIT_SYMBOL = Symbol('omit-value');
 const SERIALIZED_OMIT = "__mockttp__transform__omit__";
 
+// We play some games to preserve undefined values during serialization, because we differentiate them
+// in some transforms from null/not-present keys.
+const mapOmitToUndefined = <T extends { [key: string]: any }>(
+    input: T
+): { [K in keyof T]: T[K] | undefined } =>
+    _.mapValues(input, (v) =>
+        v === SERIALIZED_OMIT || v === OMIT_SYMBOL
+            ? undefined // Replace our omit placeholders with actual undefineds
+            : v
+    );
+
 export class PassThroughHandler extends Serializable implements RequestHandler {
     readonly type = 'passthrough';
 
@@ -1509,10 +1526,17 @@ export class PassThroughHandler extends Serializable implements RequestHandler {
             ),
             transformRequest: {
                 ...this.transformRequest,
+                // Body is always serialized as a base64 buffer:
                 replaceBody: !!this.transformRequest?.replaceBody
-                    // Always serialize as a base64 buffer:
                     ? serializeBuffer(asBuffer(this.transformRequest.replaceBody))
                     : undefined,
+                // Update objects need to capture undefined & null as distict values:
+                updateHeaders: !!this.transformRequest?.updateHeaders
+                    ? JSON.stringify(
+                        this.transformRequest.updateHeaders,
+                        (k, v) => v === undefined ? SERIALIZED_OMIT : v
+                    )
+                    : undefined,
                 updateJsonBody: !!this.transformRequest?.updateJsonBody
                     ? JSON.stringify(
                         this.transformRequest.updateJsonBody,
@@ -1522,10 +1546,17 @@ export class PassThroughHandler extends Serializable implements RequestHandler {
             },
             transformResponse: {
                 ...this.transformResponse,
+                // Body is always serialized as a base64 buffer:
                 replaceBody: !!this.transformResponse?.replaceBody
-                    // Always serialize as a base64 buffer:
                     ? serializeBuffer(asBuffer(this.transformResponse.replaceBody))
                     : undefined,
+                // Update objects need to capture undefined & null as distict values:
+                updateHeaders: !!this.transformResponse?.updateHeaders
+                    ? JSON.stringify(
+                        this.transformResponse.updateHeaders,
+                        (k, v) => v === undefined ? SERIALIZED_OMIT : v
+                    )
+                    : undefined,
                 updateJsonBody: !!this.transformResponse?.updateJsonBody
                     ? JSON.stringify(
                         this.transformResponse.updateJsonBody,
@@ -1581,25 +1612,25 @@ export class PassThroughHandler extends Serializable implements RequestHandler {
                 ...(data.transformRequest?.replaceBody !== undefined ? {
                     replaceBody: deserializeBuffer(data.transformRequest.replaceBody)
                 } : {}),
+                ...(data.transformRequest?.updateHeaders !== undefined ? {
+                    updateHeaders: mapOmitToUndefined(JSON.parse(data.transformRequest.updateHeaders))
+                } : {}),
                 ...(data.transformRequest?.updateJsonBody !== undefined ? {
-                    updateJsonBody: JSON.parse(
-                        data.transformRequest.updateJsonBody,
-                        (k, v) => v === SERIALIZED_OMIT ? OMIT_SYMBOL : v
-                    )
+                    updateJsonBody: mapOmitToUndefined(JSON.parse(data.transformRequest.updateJsonBody))
                 } : {}),
-            },
+            } as RequestTransform,
             transformResponse: {
                 ...data.transformResponse,
                 ...(data.transformResponse?.replaceBody !== undefined ? {
                     replaceBody: deserializeBuffer(data.transformResponse.replaceBody)
                 } : {}),
+                ...(data.transformResponse?.updateHeaders !== undefined ? {
+                    updateHeaders: mapOmitToUndefined(JSON.parse(data.transformResponse.updateHeaders))
+                } : {}),
                 ...(data.transformResponse?.updateJsonBody !== undefined ? {
-                    updateJsonBody: JSON.parse(
-                        data.transformResponse.updateJsonBody,
-                        (k, v) => v === SERIALIZED_OMIT ? OMIT_SYMBOL : v
-                    )
+                    updateJsonBody: mapOmitToUndefined(JSON.parse(data.transformResponse.updateJsonBody))
                 } : {})
-            },
+            } as ResponseTransform,
             // Backward compat for old clients:
             ...data.forwardToLocation ? {
                 forwarding: { targetHost: data.forwardToLocation }

test/integration/remote-client.spec.ts

@@ -192,16 +192,23 @@ nodeOnly(() => {
 
                     await client.post(targetServer.urlFor('/req')).thenPassThrough({
                         transformRequest: {
+                            updateHeaders: {
+                                'custom-header': undefined, // Remove
+                                'injected-header': 'injected-value' // Add
+                            },
                             updateJsonBody: {
-                                initialValue: undefined,
-                                replacementValue: true
+                                initialValue: undefined, // Remove
+                                replacementValue: true // Add
                             }
                         }
                     });
 
                     expect(await request.post(targetServer.urlFor('/req'), {
                         proxy: client.urlFor("/"),
                         json: true,
+                        headers: {
+                            'custom-header': 'a custom value'
+                        },
                         body: {
                             initialValue: true
                         }
@@ -213,7 +220,8 @@ nodeOnly(() => {
                             accept: 'application/json',
                             'content-type': 'application/json',
                             'content-length': '25',
-                            connection: 'close'
+                            connection: 'close',
+                            'injected-header': 'injected-value' // Only injected header remains
                         },
                         body: JSON.stringify({
                             // Initial value is removed, only replacement remains:
@@ -223,17 +231,30 @@ nodeOnly(() => {
 
                     await client.post(targetServer.urlFor('/res')).thenPassThrough({
                         transformResponse: {
+                            updateHeaders: {
+                                'custom-header': undefined, // Remove
+                                'injected-header': 'injected-value' // Add
+                            },
                             updateJsonBody: {
                                 method: 'REPLACEMENT METHOD',
                                 headers: undefined
                             }
                         }
                     });
 
-                    expect(await request.post(targetServer.urlFor('/res'), {
+                    const response = await request.post(targetServer.urlFor('/res'), {
                         proxy: client.urlFor("/"),
-                        json: true
-                    })).to.deep.equal({
+                        json: true,
+                        resolveWithFullResponse: true
+                    });
+
+                    expect(response.headers).to.deep.equal({
+                        'access-control-allow-origin': '*',
+                        'content-type': 'application/json',
+                        'injected-header': 'injected-value'
+                    });
+
+                    expect(response.body).to.deep.equal({
                         url: `http://localhost:${targetServer.port}/res`,
                         // Method field replaced, headers field removed:
                         method: 'REPLACEMENT METHOD',