[dotnet] Add transformation for network responses via CDP

SeleniumHQ · Sep 21, 2021 · 2a03bde · 2a03bde
1 parent e49e621
commit 2a03bde
Show file tree

Hide file tree

Showing 6 changed files with 144 additions and 33 deletions.
diff --git a/dotnet/src/webdriver/DevTools/v85/V85Network.cs b/dotnet/src/webdriver/DevTools/v85/V85Network.cs
@@ -289,22 +289,25 @@ private void OnFetchRequestPaused(object sender, Fetch.RequestPausedEventArgs e)
                     wrappedResponse.ResponseData.StatusCode = e.ResponseStatusCode.Value;
                 }
 
-                foreach (var header in e.ResponseHeaders)
+                if (e.ResponseHeaders != null)
                 {
-                    if (header.Name.ToLowerInvariant() == "set-cookie")
+                    foreach (var header in e.ResponseHeaders)
                     {
-                        wrappedResponse.ResponseData.CookieHeaders.Add(header.Value);
-                    }
-                    else
-                    {
-                        if (wrappedResponse.ResponseData.Headers.ContainsKey(header.Name))
+                        if (header.Name.ToLowerInvariant() == "set-cookie")
                         {
-                            string currentHeaderValue = wrappedResponse.ResponseData.Headers[header.Name];
-                            wrappedResponse.ResponseData.Headers[header.Name] = currentHeaderValue + ", " + header.Value;
+                            wrappedResponse.ResponseData.CookieHeaders.Add(header.Value);
                         }
                         else
                         {
-                            wrappedResponse.ResponseData.Headers.Add(header.Name, header.Value);
+                            if (wrappedResponse.ResponseData.Headers.ContainsKey(header.Name))
+                            {
+                                string currentHeaderValue = wrappedResponse.ResponseData.Headers[header.Name];
+                                wrappedResponse.ResponseData.Headers[header.Name] = currentHeaderValue + ", " + header.Value;
+                            }
+                            else
+                            {
+                                wrappedResponse.ResponseData.Headers.Add(header.Name, header.Value);
+                            }
                         }
                     }
                 }

diff --git a/dotnet/src/webdriver/DevTools/v92/V92Network.cs b/dotnet/src/webdriver/DevTools/v92/V92Network.cs
@@ -290,22 +290,25 @@ private void OnFetchRequestPaused(object sender, Fetch.RequestPausedEventArgs e)
                     wrappedResponse.ResponseData.StatusCode = e.ResponseStatusCode.Value;
                 }
 
-                foreach (var header in e.ResponseHeaders)
+                if (e.ResponseHeaders != null)
                 {
-                    if (header.Name.ToLowerInvariant() == "set-cookie")
+                    foreach (var header in e.ResponseHeaders)
                     {
-                        wrappedResponse.ResponseData.CookieHeaders.Add(header.Value);
-                    }
-                    else
-                    {
-                        if (wrappedResponse.ResponseData.Headers.ContainsKey(header.Name))
+                        if (header.Name.ToLowerInvariant() == "set-cookie")
                         {
-                            string currentHeaderValue = wrappedResponse.ResponseData.Headers[header.Name];
-                            wrappedResponse.ResponseData.Headers[header.Name] = currentHeaderValue + ", " + header.Value;
+                            wrappedResponse.ResponseData.CookieHeaders.Add(header.Value);
                         }
                         else
                         {
-                            wrappedResponse.ResponseData.Headers.Add(header.Name, header.Value);
+                            if (wrappedResponse.ResponseData.Headers.ContainsKey(header.Name))
+                            {
+                                string currentHeaderValue = wrappedResponse.ResponseData.Headers[header.Name];
+                                wrappedResponse.ResponseData.Headers[header.Name] = currentHeaderValue + ", " + header.Value;
+                            }
+                            else
+                            {
+                                wrappedResponse.ResponseData.Headers.Add(header.Name, header.Value);
+                            }
                         }
                     }
                 }

diff --git a/dotnet/src/webdriver/DevTools/v93/V93Network.cs b/dotnet/src/webdriver/DevTools/v93/V93Network.cs
@@ -289,22 +289,25 @@ private void OnFetchRequestPaused(object sender, Fetch.RequestPausedEventArgs e)
                     wrappedResponse.ResponseData.StatusCode = e.ResponseStatusCode.Value;
                 }
 
-                foreach (var header in e.ResponseHeaders)
+                if (e.ResponseHeaders != null)
                 {
-                    if (header.Name.ToLowerInvariant() == "set-cookie")
+                    foreach (var header in e.ResponseHeaders)
                     {
-                        wrappedResponse.ResponseData.CookieHeaders.Add(header.Value);
-                    }
-                    else
-                    {
-                        if (wrappedResponse.ResponseData.Headers.ContainsKey(header.Name))
+                        if (header.Name.ToLowerInvariant() == "set-cookie")
                         {
-                            string currentHeaderValue = wrappedResponse.ResponseData.Headers[header.Name];
-                            wrappedResponse.ResponseData.Headers[header.Name] = currentHeaderValue + ", " + header.Value;
+                            wrappedResponse.ResponseData.CookieHeaders.Add(header.Value);
                         }
                         else
                         {
-                            wrappedResponse.ResponseData.Headers.Add(header.Name, header.Value);
+                            if (wrappedResponse.ResponseData.Headers.ContainsKey(header.Name))
+                            {
+                                string currentHeaderValue = wrappedResponse.ResponseData.Headers[header.Name];
+                                wrappedResponse.ResponseData.Headers[header.Name] = currentHeaderValue + ", " + header.Value;
+                            }
+                            else
+                            {
+                                wrappedResponse.ResponseData.Headers.Add(header.Name, header.Value);
+                            }
                         }
                     }
                 }

diff --git a/dotnet/src/webdriver/INetwork.cs b/dotnet/src/webdriver/INetwork.cs
@@ -63,6 +63,18 @@ public interface INetwork
         /// </summary>
         void ClearAuthenticationHandlers();
 
+        /// <summary>
+        /// Adds a <see cref="NetworkResponseHandler"/> to examine received network responses,
+        /// and optionally modify the response. 
+        /// </summary>
+        /// <param name="handler">The <see cref="NetworkResponseHandler"/> to add.</param>
+        void AddResponseHandler(NetworkResponseHandler handler);
+
+        /// <summary>
+        /// Clears all added <see cref="NetworkResponseHandler"/> instances.
+        /// </summary>
+        void ClearResponseHandlers();
+
         /// <summary>
         /// Asynchronously starts monitoring for network traffic.
         /// </summary>

diff --git a/dotnet/src/webdriver/NetworkManager.cs b/dotnet/src/webdriver/NetworkManager.cs
@@ -30,10 +30,11 @@ public class NetworkManager : INetwork
     {
         private Lazy<DevToolsSession> session;
         private List<NetworkRequestHandler> requestHandlers = new List<NetworkRequestHandler>();
+        private List<NetworkResponseHandler> responseHandlers = new List<NetworkResponseHandler>();
         private List<NetworkAuthenticationHandler> authenticationHandlers = new List<NetworkAuthenticationHandler>();
 
         /// <summary>
-        /// Initializes a new instance of the <see cref="RemoteNetwork"/> class.
+        /// Initializes a new instance of the <see cref="NetworkManager"/> class.
         /// </summary>
         /// <param name="driver">The <see cref="IWebDriver"/> instance on which the network should be monitored.</param>
         public NetworkManager(IWebDriver driver)
@@ -161,6 +162,34 @@ public void ClearAuthenticationHandlers()
             this.authenticationHandlers.Clear();
         }
 
+        /// <summary>
+        /// Adds a <see cref="NetworkResponseHandler"/> to examine received network responses,
+        /// and optionally modify the response. 
+        /// </summary>
+        /// <param name="handler">The <see cref="NetworkResponseHandler"/> to add.</param>
+        public void AddResponseHandler(NetworkResponseHandler handler)
+        {
+            if (handler == null)
+            {
+                throw new ArgumentNullException("handler", "Request handler cannot be null");
+            }
+
+            if (handler.ResponseMatcher == null)
+            {
+                throw new ArgumentException("Matcher for response cannot be null", "handler");
+            }
+
+            this.responseHandlers.Add(handler);
+        }
+
+        /// <summary>
+        /// Clears all added <see cref="NetworkResponseHandler"/> instances.
+        /// </summary>
+        public void ClearResponseHandlers()
+        {
+            this.responseHandlers.Clear();
+        }
+
         private async void OnAuthRequired(object sender, AuthRequiredEventArgs e)
         {
             string requestId = e.RequestId;
@@ -213,12 +242,32 @@ private async void OnRequestPaused(object sender, RequestPausedEventArgs e)
 
         private async void OnResponsePaused(object sender, ResponsePausedEventArgs e)
         {
-            await this.session.Value.Domains.Network.AddResponseBody(e.ResponseData);
-            await this.session.Value.Domains.Network.ContinueResponseWithoutModification(e.ResponseData);
+            if (e.ResponseData.Headers.Count > 0)
+            {
+                // If no headers are present, the body cannot be retrieved.
+                await this.session.Value.Domains.Network.AddResponseBody(e.ResponseData);
+            }
+
             if (this.NetworkResponseReceived != null)
             {
                 this.NetworkResponseReceived(this, new NetworkResponseReceivedEventArgs(e.ResponseData));
             }
+
+            foreach (var handler in this.responseHandlers)
+            {
+                if (handler.ResponseMatcher.Invoke(e.ResponseData))
+                {
+                    // NOTE: We create a dummy HttpRequestData object here, because the ContinueRequestWithResponse
+                    // method demands one; however, the only property used by that method is the RequestId property.
+                    // It might be better to refactor that method signature to simply pass the request ID, or
+                    // alternatively, just pass the response data, which should also contain the request ID anyway.
+                    HttpRequestData requestData = new HttpRequestData() { RequestId = e.ResponseData.RequestId };
+                    await this.session.Value.Domains.Network.ContinueRequestWithResponse(requestData, handler.ResponseTransformer(e.ResponseData));
+                    return;
+                }
+            }
+
+            await this.session.Value.Domains.Network.ContinueResponseWithoutModification(e.ResponseData);
         }
     }
 }
diff --git a/dotnet/src/webdriver/NetworkResponseHandler.cs b/dotnet/src/webdriver/NetworkResponseHandler.cs
@@ -0,0 +1,41 @@
+// <copyright file="NetworkResponsetHandler.cs" company="WebDriver Committers">
+// Licensed to the Software Freedom Conservancy (SFC) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The SFC licenses this file
+// to you under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// </copyright>
+
+using System;
+
+namespace OpenQA.Selenium
+{
+    /// <summary>
+    /// Allows a user to handle a returned network, potentially modifying it before processing by the browser.
+    /// </summary>
+    public class NetworkResponseHandler
+    {
+        /// <summary>
+        /// Gets or sets a function that evaluates returned response data in an <see cref="HttpResponseData"/> object,
+        /// and returns a value indicating whether the data matches the specified criteria.
+        /// </summary>
+        public Func<HttpResponseData, bool> ResponseMatcher { get; set; }
+
+        /// <summary>
+        /// Gets or sets a function that accepts an <see cref="HttpResponseData"/> object describing a network
+        /// response received by the browser, and returns a modified <see cref="HttpResponseData"/> object to used
+        /// as the actual network response.
+        /// </summary>
+        public Func<HttpResponseData, HttpResponseData> ResponseTransformer { get; set; }
+    }
+}