From ce0b6f1c80be53f8e69b72cec2b67258d4b87fa6 Mon Sep 17 00:00:00 2001
From: Puja Jagani <puja.jagani93@gmail.com>
Date: Thu, 2 Feb 2023 13:06:55 +0530
Subject: [PATCH] [java] [cdp] Simplify Augmentation for Basic Auth (#11601)

---
 .../selenium/remote/AddHasAuthentication.java | 68 +++++++++++++++++++
 .../org/openqa/selenium/remote/Augmenter.java | 32 +++++++++
 2 files changed, 100 insertions(+)
 create mode 100644 java/src/org/openqa/selenium/remote/AddHasAuthentication.java

diff --git a/java/src/org/openqa/selenium/remote/AddHasAuthentication.java b/java/src/org/openqa/selenium/remote/AddHasAuthentication.java
new file mode 100644
index 0000000000000..e4db2a602845a
--- /dev/null
+++ b/java/src/org/openqa/selenium/remote/AddHasAuthentication.java
@@ -0,0 +1,68 @@
+// 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.
+
+package org.openqa.selenium.remote;
+
+import static org.openqa.selenium.remote.Browser.CHROME;
+import static org.openqa.selenium.remote.Browser.EDGE;
+import static org.openqa.selenium.remote.Browser.OPERA;
+
+import org.openqa.selenium.Capabilities;
+import org.openqa.selenium.HasAuthentication;
+import org.openqa.selenium.WebDriver;
+import org.openqa.selenium.devtools.DevTools;
+import org.openqa.selenium.devtools.HasDevTools;
+import org.openqa.selenium.internal.Require;
+
+import java.util.function.Predicate;
+import java.util.logging.Logger;
+
+public class AddHasAuthentication
+  implements AugmenterProvider<HasAuthentication> {
+
+  private static final Logger logger = Logger.getLogger(AddHasAuthentication.class.getName());
+  private static final Predicate<String> IS_CHROMIUM_BROWSER = name ->
+    CHROME.is(name) ||
+    EDGE.is(name) ||
+    OPERA.is(name);
+
+  @Override
+  public Predicate<Capabilities> isApplicable() {
+    return caps -> IS_CHROMIUM_BROWSER.test(caps.getBrowserName());
+  }
+
+  @Override
+  public Class<HasAuthentication> getDescribedInterface() {
+    return HasAuthentication.class;
+  }
+
+  @Override
+  public HasAuthentication getImplementation(Capabilities capabilities, ExecuteMethod executeMethod) {
+    return (whenThisMatches, useTheseCredentials) -> {
+      Require.nonNull("Check to use to see how we should authenticate", whenThisMatches);
+      Require.nonNull("Credentials to use when authenticating", useTheseCredentials);
+
+      if (((RemoteExecuteMethod) executeMethod).getWrappedDriver() instanceof HasDevTools) {
+        WebDriver driver = ((RemoteExecuteMethod) executeMethod).getWrappedDriver();
+        DevTools devTools = ((HasDevTools) driver).getDevTools();
+        devTools.createSessionIfThereIsNotOne();
+        devTools.getDomains().network().addAuthHandler(whenThisMatches, useTheseCredentials);
+      }
+      // Todo: Similarly add for BiDi once BiDi supports the same functionality
+    };
+  }
+}
diff --git a/java/src/org/openqa/selenium/remote/Augmenter.java b/java/src/org/openqa/selenium/remote/Augmenter.java
index 98f62087391a9..2148524d48490 100644
--- a/java/src/org/openqa/selenium/remote/Augmenter.java
+++ b/java/src/org/openqa/selenium/remote/Augmenter.java
@@ -26,6 +26,7 @@
 
 import org.openqa.selenium.Beta;
 import org.openqa.selenium.Capabilities;
+import org.openqa.selenium.HasAuthentication;
 import org.openqa.selenium.HasCapabilities;
 import org.openqa.selenium.ImmutableCapabilities;
 import org.openqa.selenium.WebDriver;
@@ -36,6 +37,7 @@
 
 import java.lang.reflect.Field;
 import java.lang.reflect.Modifier;
+import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
 import java.util.ServiceLoader;
@@ -158,6 +160,12 @@ public WebDriver augment(WebDriver driver) {
       .filter(augmentation -> augmentation.whenMatches.test(caps))
       .collect(Collectors.toList());
 
+    return this.augment(driver, matchingAugmenters);
+  }
+
+  private WebDriver augment(WebDriver driver, List<Augmentation<?>> matchingAugmenters) {
+    Capabilities caps = ImmutableCapabilities.copyOf(((HasCapabilities) driver).getCapabilities());
+
     if (matchingAugmenters.isEmpty()) {
       return driver;
     }
@@ -193,12 +201,36 @@ public WebDriver augment(WebDriver driver) {
 
       copyFields(driver.getClass(), driver, toReturn);
 
+      toReturn = addDependentAugmentations(toReturn);
+
       return toReturn;
     } catch (ReflectiveOperationException e) {
       throw new IllegalStateException("Unable to create new proxy", e);
     }
   }
 
+  private WebDriver addDependentAugmentations(WebDriver driver) {
+    List<Augmentation<?>> augmentationList = new ArrayList<>();
+
+    WebDriver toReturn = driver;
+
+    // add interfaces that need to use the augmented driver
+    if (!(driver instanceof HasAuthentication)) {
+      augmentationList.add(createAugmentation(new AddHasAuthentication()));
+    }
+
+    if (!augmentationList.isEmpty()) {
+      Capabilities caps = ImmutableCapabilities.copyOf(((HasCapabilities) driver).getCapabilities());
+
+      List<Augmentation<?>> matchingAugmenters = augmentationList.stream()
+        .filter(augmentation -> augmentation.whenMatches.test(caps))
+        .collect(Collectors.toList());
+
+      toReturn = this.augment(driver, matchingAugmenters);
+    }
+    return toReturn;
+  }
+
   private RemoteWebDriver extractRemoteWebDriver(WebDriver driver) {
     Require.nonNull("WebDriver", driver);