improve prevent-xhr - add mock getResponseHeader() and getAllResponseHeaders() methods

slavaleleka · slavaleleka · commit b378dd4b49a5 · 2023-05-12T13:49:31.000+03:00
diff --git a/src/scriptlets/prevent-xhr.js b/src/scriptlets/prevent-xhr.js
@@ -3,6 +3,7 @@ import {
     objectToString,
     generateRandomResponse,
     matchRequestProps,
+    getXhrData,
     logMessage,
     // following helpers should be imported and injected
     // because they are used by helpers above
@@ -95,16 +96,18 @@ export function preventXHR(source, propsToMatch, customResponseText) {
         return;
     }
 
-    let response = '';
-    let responseText = '';
-    let responseUrl;
+    const nativeOpen = window.XMLHttpRequest.prototype.open;
+    const nativeSend = window.XMLHttpRequest.prototype.send;
+
+    let xhrData;
+    let modifiedResponse = '';
+    let modifiedResponseText = '';
+
     const openWrapper = (target, thisArg, args) => {
-        // Get method and url from .open()
-        const xhrData = {
-            method: args[0],
-            url: args[1],
-        };
-        responseUrl = xhrData.url;
+        // Get original request properties
+        // eslint-disable-next-line prefer-spread
+        xhrData = getXhrData.apply(null, args);
+
         if (typeof propsToMatch === 'undefined') {
             // Log if no propsToMatch given
             logMessage(source, `xhr( ${objectToString(xhrData)} )`, true);
@@ -113,6 +116,22 @@ export function preventXHR(source, propsToMatch, customResponseText) {
             thisArg.shouldBePrevented = true;
         }
 
+        // Trap setRequestHeader of target xhr object to mimic request headers later;
+        // needed for getResponseHeader() and getAllResponseHeaders() methods
+        if (thisArg.shouldBePrevented) {
+            thisArg.collectedHeaders = [];
+            const setRequestHeaderWrapper = (target, thisArg, args) => {
+                // Collect headers
+                thisArg.collectedHeaders.push(args);
+                return Reflect.apply(target, thisArg, args);
+            };
+            const setRequestHeaderHandler = {
+                apply: setRequestHeaderWrapper,
+            };
+            // setRequestHeader() can only be called on xhr.open(),
+            // so we can safely proxy it here
+            thisArg.setRequestHeader = new Proxy(thisArg.setRequestHeader, setRequestHeaderHandler);
+        }
         return Reflect.apply(target, thisArg, args);
     };
 
@@ -122,57 +141,143 @@ export function preventXHR(source, propsToMatch, customResponseText) {
         }
 
         if (thisArg.responseType === 'blob') {
-            response = new Blob();
+            modifiedResponse = new Blob();
         }
-
         if (thisArg.responseType === 'arraybuffer') {
-            response = new ArrayBuffer();
+            modifiedResponse = new ArrayBuffer();
         }
 
         if (customResponseText) {
             const randomText = generateRandomResponse(customResponseText);
             if (randomText) {
-                responseText = randomText;
+                modifiedResponseText = randomText;
             } else {
+                // FIXME: improve error text
                 logMessage(source, `Invalid range: ${customResponseText}`);
             }
         }
-        // Mock response object
-        Object.defineProperties(thisArg, {
-            readyState: { value: 4, writable: false },
-            response: { value: response, writable: false },
-            responseText: { value: responseText, writable: false },
-            responseURL: { value: responseUrl, writable: false },
-            responseXML: { value: '', writable: false },
-            status: { value: 200, writable: false },
-            statusText: { value: 'OK', writable: false },
+
+        /**
+         * Create separate XHR request with original request's input
+         * to be able to collect response data without triggering
+         * listeners on original XHR object
+         */
+        const forgedRequest = new XMLHttpRequest();
+        forgedRequest.addEventListener('readystatechange', () => {
+            if (forgedRequest.readyState !== 4) {
+                return;
+            }
+
+            const {
+                readyState,
+                responseURL,
+                responseXML,
+                status,
+                statusText,
+            } = forgedRequest;
+
+            // Mock response object
+            Object.defineProperties(thisArg, {
+                // original values
+                readyState: { value: readyState, writable: false },
+                status: { value: status, writable: false },
+                statusText: { value: statusText, writable: false },
+                responseURL: { value: responseURL, writable: false },
+                responseXML: { value: responseXML, writable: false },
+                // modified values
+                response: { value: modifiedResponse, writable: false },
+                responseText: { value: modifiedResponseText, writable: false },
+            });
+
+            // Mock events
+            setTimeout(() => {
+                const stateEvent = new Event('readystatechange');
+                thisArg.dispatchEvent(stateEvent);
+
+                const loadEvent = new Event('load');
+                thisArg.dispatchEvent(loadEvent);
+
+                const loadEndEvent = new Event('loadend');
+                thisArg.dispatchEvent(loadEndEvent);
+            }, 1);
+
+            hit(source);
         });
-        // Mock events
-        setTimeout(() => {
-            const stateEvent = new Event('readystatechange');
-            thisArg.dispatchEvent(stateEvent);
 
-            const loadEvent = new Event('load');
-            thisArg.dispatchEvent(loadEvent);
+        nativeOpen.apply(forgedRequest, [xhrData.method, xhrData.url]);
+
+        // Mimic request headers before sending
+        // setRequestHeader can only be called on open request objects
+        thisArg.collectedHeaders.forEach((header) => {
+            const name = header[0];
+            const value = header[1];
+            forgedRequest.setRequestHeader(name, value);
+        });
 
-            const loadEndEvent = new Event('loadend');
-            thisArg.dispatchEvent(loadEndEvent);
-        }, 1);
+        try {
+            nativeSend.call(forgedRequest, args);
+        } catch {
+            return Reflect.apply(target, thisArg, args);
+        }
 
-        hit(source);
         return undefined;
     };
 
+    const getHeaderWrapper = (target, thisArg, args) => {
+        if (!thisArg.collectedHeaders.length) {
+            return null;
+        }
+        // The search for the header name is case-insensitive
+        // https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/getResponseHeader
+        const searchHeaderName = args[0].toLowerCase();
+        const matchedHeader = thisArg.collectedHeaders.find((header) => {
+            const headerName = header[0].toLowerCase();
+            return headerName === searchHeaderName;
+        });
+        return matchedHeader
+            ? matchedHeader[1]
+            : null;
+    };
+
+    const getAllHeadersWrapper = (target, thisArg) => {
+        if (!thisArg.collectedHeaders.length) {
+            return '';
+        }
+        const allHeadersStr = thisArg.collectedHeaders
+            .map((header) => {
+                const headerName = header[0];
+                const headerValue = header[1];
+                // In modern browsers, the header names are returned in all lower case, as per the latest spec.
+                // https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/getAllResponseHeaders
+                return `${headerName.toLowerCase()}: ${headerValue}`;
+            })
+            .join('\r\n');
+        return allHeadersStr;
+    };
+
     const openHandler = {
         apply: openWrapper,
     };
-
     const sendHandler = {
         apply: sendWrapper,
     };
+    const getHeaderHandler = {
+        apply: getHeaderWrapper,
+    };
+    const getAllHeadersHandler = {
+        apply: getAllHeadersWrapper,
+    };
 
     XMLHttpRequest.prototype.open = new Proxy(XMLHttpRequest.prototype.open, openHandler);
     XMLHttpRequest.prototype.send = new Proxy(XMLHttpRequest.prototype.send, sendHandler);
+    XMLHttpRequest.prototype.getResponseHeader = new Proxy(
+        XMLHttpRequest.prototype.getResponseHeader,
+        getHeaderHandler,
+    );
+    XMLHttpRequest.prototype.getAllResponseHeaders = new Proxy(
+        XMLHttpRequest.prototype.getAllResponseHeaders,
+        getAllHeadersHandler,
+    );
 }
 
 preventXHR.names = [
@@ -185,10 +290,11 @@ preventXHR.names = [
 
 preventXHR.injections = [
     hit,
-    logMessage,
     objectToString,
-    matchRequestProps,
     generateRandomResponse,
+    matchRequestProps,
+    getXhrData,
+    logMessage,
     toRegExp,
     isValidStrPattern,
     escapeRegExp,
diff --git a/tests/scriptlets/prevent-xhr.test.js b/tests/scriptlets/prevent-xhr.test.js
@@ -57,7 +57,8 @@ if (isSupported) {
             if (input.indexOf('trace') > -1) {
                 return;
             }
-            const EXPECTED_LOG_STR = `${name}: xhr( method:"${METHOD}" url:"${URL}" )`;
+            // eslint-disable-next-line max-len
+            const EXPECTED_LOG_STR = `${name}: xhr( method:"${METHOD}" url:"${URL}" async:"undefined" user:"undefined" password:"undefined" )`;
             assert.ok(startsWith(input, EXPECTED_LOG_STR), 'console.hit input');
         };
 
@@ -96,6 +97,56 @@ if (isSupported) {
         xhr.send();
     });
 
+    test('Empty arg to prevent all, check getResponseHeader() and getAllResponseHeaders() methods', async (assert) => {
+        const METHOD = 'GET';
+        const URL = `${FETCH_OBJECTS_PATH}/test01.json`;
+        const MATCH_DATA = [''];
+        const HEADER_NAME_1 = 'Test-Type';
+        const HEADER_VALUE_1 = 'application/json';
+        const HEADER_NAME_2 = 'Test-Length';
+        const HEADER_VALUE_2 = '12345';
+        const ABSENT_HEADER_NAME = 'Test-Absent';
+
+        runScriptlet(name, MATCH_DATA);
+
+        const done = assert.async();
+
+        const xhr = new XMLHttpRequest();
+        xhr.open(METHOD, URL);
+        xhr.setRequestHeader(HEADER_NAME_1, HEADER_VALUE_1);
+        xhr.setRequestHeader(HEADER_NAME_2, HEADER_VALUE_2);
+
+        xhr.onload = () => {
+            assert.strictEqual(xhr.readyState, 4, 'Response done');
+            assert.strictEqual(xhr.response, '', 'Response data mocked');
+            assert.strictEqual(window.hit, 'FIRED', 'hit function fired');
+            done();
+        };
+        xhr.send();
+
+        assert.strictEqual(
+            xhr.getResponseHeader(HEADER_NAME_1),
+            HEADER_VALUE_1,
+            'getResponseHeader() is mocked, value 1 returned',
+        );
+        assert.strictEqual(
+            xhr.getResponseHeader(HEADER_NAME_2),
+            HEADER_VALUE_2,
+            'getResponseHeader() is mocked',
+        );
+        assert.strictEqual(
+            xhr.getResponseHeader(ABSENT_HEADER_NAME),
+            null,
+            'getResponseHeader() is mocked, null returned for non-existent header',
+        );
+
+        const expectedAllHeaders = [
+            `${HEADER_NAME_1.toLowerCase()}: ${HEADER_VALUE_1}`,
+            `${HEADER_NAME_2.toLowerCase()}: ${HEADER_VALUE_2}`,
+        ].join('\r\n');
+        assert.strictEqual(xhr.getAllResponseHeaders(), expectedAllHeaders, 'getAllResponseHeaders() is mocked');
+    });
+
     test('Empty arg, prevent all, randomize response text', async (assert) => {
         const METHOD = 'GET';
         const URL = `${FETCH_OBJECTS_PATH}/test01.json`;
