Merge pull request #186 from fitudao3788/main

Implement req in beforeResponse #80

Author: pimterry

src/rules/requests/request-handler-definitions.ts

@@ -550,17 +550,19 @@ export interface PassThroughHandlerOptions extends PassThroughHandlerConnectionO
     /**
      * A callback that will be passed the full response before it is passed through,
      * and which returns a value that defines how the the response content should
-     * be transformed before it's returned to the client.
+     * be transformed before it's returned to the client. The callback is also passed
+     * the request that was sent to the server (as a 2nd parameter) for reference.
      *
      * The callback can either return an object to define how the response should be
-     * changed, or the string 'close' to immediately close the underlying connection.
+     * changed, or the strings 'close' or 'reset' to immediately close/reset the
+     * underlying connection.
      *
      * All fields on the object are optional, and returning undefined is equivalent
      * to returning an empty object (transforming nothing).
      *
      * See {@link CallbackResponseMessageResult} for the possible fields that can be set.
      */
-    beforeResponse?: (res: PassThroughResponse) => MaybePromise<CallbackResponseResult | void> | void;
+    beforeResponse?: (res: PassThroughResponse, req: CompletedRequest) => MaybePromise<CallbackResponseResult | void> | void;
 }
 
 export interface RequestTransform {
@@ -754,7 +756,7 @@ export interface BeforePassthroughRequestRequest {
  * @internal
  */
 export interface BeforePassthroughResponseRequest {
-    args: [Replace<PassThroughResponse, { body: string }>];
+    args: [Replace<PassThroughResponse, { body: string }>, Replace<CompletedRequest, { body: string }>];
 }
 
 /**
@@ -780,7 +782,7 @@ export class PassThroughHandlerDefinition extends Serializable implements Reques
 
     public readonly beforeRequest?: (req: CompletedRequest) =>
         MaybePromise<CallbackRequestResult | void> | void;
-    public readonly beforeResponse?: (res: PassThroughResponse) =>
+    public readonly beforeResponse?: (res: PassThroughResponse, req: CompletedRequest) =>
         MaybePromise<CallbackResponseResult | void> | void;
 
     public readonly proxyConfig?: ProxyConfig;
@@ -925,7 +927,8 @@ export class PassThroughHandlerDefinition extends Serializable implements Reques
                 CallbackResponseResult | undefined
             >('beforeResponse', async (req) => {
                 const callbackResult = await this.beforeResponse!(
-                    withDeserializedBodyReader(req.args[0])
+                    withDeserializedBodyReader(req.args[0]),
+                    withDeserializedBodyReader(req.args[1]),
                 );
 
                 if (typeof callbackResult === 'string') {

src/rules/requests/request-handlers.ts

@@ -924,13 +924,27 @@ export class PassThroughHandler extends PassThroughHandlerDefinition {
                     originalBody = await streamToBuffer(serverRes);
                     let serverHeaders = rawHeadersToObject(serverRawHeaders);
 
+                    let reqHeader = rawHeadersToObjectPreservingCase(rawHeaders);
                     modifiedRes = await this.beforeResponse({
                         id: clientReq.id,
                         statusCode: serverStatusCode,
                         statusMessage: serverRes.statusMessage,
                         headers: serverHeaders,
                         rawHeaders: _.cloneDeep(serverRawHeaders),
                         body: buildBodyReader(originalBody, serverHeaders)
+                    }, {
+                        id: clientReq.id,
+                        protocol: protocol?.replace(':', '') ?? '',
+                        method: method,
+                        url: reqUrl,
+                        path: path ?? '',
+                        headers: reqHeader,
+                        rawHeaders: rawHeaders,
+                        timingEvents: clientReq.timingEvents,
+                        tags: clientReq.tags,
+                        body: buildBodyReader(reqBodyOverride ? Buffer.from(reqBodyOverride.buffer) : await clientReq.body.asDecodedBuffer(), reqHeader),
+                        rawTrailers: clientReq.rawTrailers ?? [],
+                        trailers: rawHeadersToObject(clientReq.rawTrailers ?? []),
                     });
 
                     if (modifiedRes === 'close') {
@@ -1214,17 +1228,17 @@ export class PassThroughHandler extends PassThroughHandlerDefinition {
             };
         }
 
-        let beforeResponse: ((res: PassThroughResponse) => MaybePromise<CallbackResponseResult | void>) | undefined;
+        let beforeResponse: ((res: PassThroughResponse, req: CompletedRequest) => MaybePromise<CallbackResponseResult | void>) | undefined;
         if (data.hasBeforeResponseCallback) {
-            beforeResponse = async (res: PassThroughResponse) => {
+            beforeResponse = async (res: PassThroughResponse, req: CompletedRequest) => {
                 const callbackResult = await channel.request<
                     BeforePassthroughResponseRequest,
                     | WithSerializedCallbackBuffers<CallbackResponseMessageResult>
                     | 'close'
                     | 'reset'
                     | undefined
                 >('beforeResponse', {
-                    args: [withSerializedBodyReader(res)]
+                    args: [withSerializedBodyReader(res), withSerializedBodyReader(req)]
                 })
 
                 if (callbackResult && typeof callbackResult !== 'string') {

test/integration/proxying/http-proxying.spec.ts

@@ -778,6 +778,25 @@ nodeOnly(() => {
                 expect(response).to.equal('all good');
             });
 
+            it("should be able to examine the request in beforeResponse", async () => {
+                await remoteServer.forGet('/').thenCallback(() => ({
+                    status: 500,
+                    headers: {
+                        'UPPERCASE-HEADER': 'VALUE'
+                    }
+                }));
+
+                await server.forGet(remoteServer.urlFor("/")).thenPassThrough({
+                    beforeResponse: (_res, req) => {
+                        expect(req.url).to.equal(remoteServer.urlFor('/'));
+                        return { status: 200, body: 'got correct req url' };
+                    }
+                });
+
+                let response = await request.get(remoteServer.urlFor("/"));
+                expect(response).to.equal('got correct req url');
+            });
+
             it("should be able to rewrite a response's status", async () => {
                 await remoteServer.forGet('/').thenReply(404);
                 await server.forGet(remoteServer.urlFor("/")).thenPassThrough({