From bef24c2b0e3fa32f02312ee056ca14fccb189334 Mon Sep 17 00:00:00 2001
From: Gabriel Luong <gabriel.luong@gmail.com>
Date: Tue, 20 Feb 2024 19:27:37 -0500
Subject: [PATCH] Bug 1881128 - Add additional headers for Google requests

---
 .../fenix/components/UrlRequestInterceptor.kt |  6 +--
 .../components/UrlRequestInterceptorTest.kt   | 42 +++++++++++++++----
 2 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/app/src/main/java/org/mozilla/fenix/components/UrlRequestInterceptor.kt b/app/src/main/java/org/mozilla/fenix/components/UrlRequestInterceptor.kt
index d267a6043..74e270f70 100644
--- a/app/src/main/java/org/mozilla/fenix/components/UrlRequestInterceptor.kt
+++ b/app/src/main/java/org/mozilla/fenix/components/UrlRequestInterceptor.kt
@@ -19,8 +19,8 @@ import mozilla.components.concept.engine.request.RequestInterceptor
  */
 class UrlRequestInterceptor(private val isDeviceRamAboveThreshold: Boolean) : RequestInterceptor {
 
-    private val isGoogleSearchRequest by lazy {
-        Regex("^https://www\\.google\\.(?:.+)/search")
+    private val isGoogleRequest by lazy {
+        Regex("^https://www\\.google\\..+")
     }
 
     @VisibleForTesting
@@ -41,7 +41,7 @@ class UrlRequestInterceptor(private val isDeviceRamAboveThreshold: Boolean) : Re
         uri: String,
         isSubframeRequest: Boolean,
     ): Boolean {
-        return !isSubframeRequest && isGoogleSearchRequest.containsMatchIn(uri)
+        return !isSubframeRequest && isGoogleRequest.containsMatchIn(uri)
     }
 
     override fun onLoadRequest(
diff --git a/app/src/test/java/org/mozilla/fenix/components/UrlRequestInterceptorTest.kt b/app/src/test/java/org/mozilla/fenix/components/UrlRequestInterceptorTest.kt
index 058c47e24..61edbff0a 100644
--- a/app/src/test/java/org/mozilla/fenix/components/UrlRequestInterceptorTest.kt
+++ b/app/src/test/java/org/mozilla/fenix/components/UrlRequestInterceptorTest.kt
@@ -59,15 +59,21 @@ class UrlRequestInterceptorTest {
     fun `WHEN should intercept request is called THEN return the correct boolean value`() {
         val urlRequestInterceptor = getUrlRequestInterceptor()
 
-        assertFalse(
+        assertTrue(
             urlRequestInterceptor.shouldInterceptRequest(
-                uri = "https://getpocket.com",
+                uri = "https://www.google.com",
                 isSubframeRequest = false,
             ),
         )
-        assertFalse(
+        assertTrue(
             urlRequestInterceptor.shouldInterceptRequest(
-                uri = "https://www.google.com",
+                uri = "https://www.google.com/webhp",
+                isSubframeRequest = false,
+            ),
+        )
+        assertTrue(
+            urlRequestInterceptor.shouldInterceptRequest(
+                uri = "https://www.google.com/preferences",
                 isSubframeRequest = false,
             ),
         )
@@ -89,6 +95,13 @@ class UrlRequestInterceptorTest {
                 isSubframeRequest = false,
             ),
         )
+
+        assertFalse(
+            urlRequestInterceptor.shouldInterceptRequest(
+                uri = "https://getpocket.com",
+                isSubframeRequest = false,
+            ),
+        )
         assertFalse(
             urlRequestInterceptor.shouldInterceptRequest(
                 uri = "https://www.google.com/search?q=blue",
@@ -114,13 +127,24 @@ class UrlRequestInterceptorTest {
     }
 
     @Test
-    fun `WHEN a Google request is loaded THEN request is not intercepted`() {
+    fun `WHEN a Google request is loaded THEN request is intercepted`() {
         val uri = "https://www.google.com"
-        val response = getUrlRequestInterceptor().onLoadRequest(
-            uri = uri,
-        )
 
-        assertNull(response)
+        assertEquals(
+            RequestInterceptor.InterceptionResponse.Url(
+                url = uri,
+                flags = LoadUrlFlags.select(
+                    LOAD_FLAGS_BYPASS_LOAD_URI_DELEGATE,
+                    ALLOW_ADDITIONAL_HEADERS,
+                ),
+                additionalHeaders = mapOf(
+                    "X-Search-Subdivision" to "0",
+                ),
+            ),
+            getUrlRequestInterceptor().onLoadRequest(
+                uri = uri,
+            ),
+        )
     }
 
     @Test