From 0e4b2254c93ac5a2a00753a76b0d84e21ffc3455 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B6rg=20Sautter?= <joerg.sautter@gmx.net>
Date: Sat, 25 Mar 2023 21:40:14 +0100
Subject: [PATCH] Handle redirects inside the JdkHttpClient

---
 .../remote/http/jdk/JdkHttpClient.java        | 42 +++++++++++++++++--
 .../remote/http/jdk/JdkHttpMessages.java      | 16 ++++---
 2 files changed, 45 insertions(+), 13 deletions(-)

diff --git a/java/src/org/openqa/selenium/remote/http/jdk/JdkHttpClient.java b/java/src/org/openqa/selenium/remote/http/jdk/JdkHttpClient.java
index 47e378a07dc28..048ccd38c6ac3 100644
--- a/java/src/org/openqa/selenium/remote/http/jdk/JdkHttpClient.java
+++ b/java/src/org/openqa/selenium/remote/http/jdk/JdkHttpClient.java
@@ -39,6 +39,7 @@
 import java.io.UncheckedIOException;
 import java.net.Authenticator;
 import java.net.PasswordAuthentication;
+import java.net.ProtocolException;
 import java.net.Proxy;
 import java.net.ProxySelector;
 import java.net.SocketAddress;
@@ -62,8 +63,6 @@
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
-import static java.net.http.HttpClient.Redirect.ALWAYS;
-
 public class JdkHttpClient implements HttpClient {
   public static final Logger LOG = Logger.getLogger(JdkHttpClient.class.getName());
   private final JdkHttpMessages messages;
@@ -82,7 +81,7 @@ public class JdkHttpClient implements HttpClient {
 
     java.net.http.HttpClient.Builder builder = java.net.http.HttpClient.newBuilder()
       .connectTimeout(config.connectionTimeout())
-      .followRedirects(ALWAYS)
+      .followRedirects(java.net.http.HttpClient.Redirect.NEVER)
       .executor(executorService);
 
     Credentials credentials = config.credentials();
@@ -289,7 +288,42 @@ public HttpResponse execute(HttpRequest req) throws UncheckedIOException {
 
     BodyHandler<byte[]> byteHandler = BodyHandlers.ofByteArray();
     try {
-      return messages.createResponse(client.send(messages.createRequest(req), byteHandler));
+      URI rawUri = messages.getRawUri(req);
+
+      // We need a custom handling of redirects to:
+      // - increase the maximum number of retries to 100
+      // - avoid a downgrade of POST requests, see the javadoc of j.n.h.HttpClient.Redirect
+      // - not run into https://bugs.openjdk.org/browse/JDK-8304701
+      for (int i = 0; i < 100; i++) {
+        java.net.http.HttpRequest request = messages.createRequest(req, rawUri);
+        java.net.http.HttpResponse<byte[]> response = client.send(request, byteHandler);
+
+        switch (response.statusCode()) {
+          case 301:
+          case 302:
+          case 303:
+          case 307:
+          case 308:
+            URI location = rawUri.resolve(response.headers()
+              .firstValue("location")
+              .orElseThrow(() -> new ProtocolException(
+                "HTTP " + response.statusCode() + " without 'location' header set"
+              )));
+
+            if ("https".equalsIgnoreCase(rawUri.getScheme()) && !"https".equalsIgnoreCase(location.getScheme()) ) {
+              throw new SecurityException("Downgrade from secure to insecure connection.");
+            } else if ("wss".equalsIgnoreCase(rawUri.getScheme()) && !"wss".equalsIgnoreCase(location.getScheme()) ) {
+              throw new SecurityException("Downgrade from secure to insecure connection.");
+            }
+
+            rawUri = location;
+            continue;
+          default:
+            return messages.createResponse(response);
+        }
+      }
+
+      throw new ProtocolException("Too many redirects: 101");
     } catch (HttpTimeoutException e) {
       throw new TimeoutException(e);
     } catch (IOException e) {
diff --git a/java/src/org/openqa/selenium/remote/http/jdk/JdkHttpMessages.java b/java/src/org/openqa/selenium/remote/http/jdk/JdkHttpMessages.java
index 1cc740a240e8d..4d5c39b4ea269 100644
--- a/java/src/org/openqa/selenium/remote/http/jdk/JdkHttpMessages.java
+++ b/java/src/org/openqa/selenium/remote/http/jdk/JdkHttpMessages.java
@@ -45,8 +45,8 @@ public JdkHttpMessages(ClientConfig config) {
     this.config = Objects.requireNonNull(config, "Client config");
   }
 
-  public java.net.http.HttpRequest createRequest(HttpRequest req) {
-    String rawUrl = getRawUrl(config.baseUri(), req.getUri());
+  public java.net.http.HttpRequest createRequest(HttpRequest req, URI rawUri) {
+    String rawUrl = rawUri.toString();
 
     // Add query string if necessary
     String queryString = StreamSupport.stream(req.getQueryParameterNames().spliterator(), false)
@@ -129,20 +129,18 @@ private BodyPublisher notChunkingBodyPublisher(HttpRequest req) {
     return BodyPublishers.fromPublisher(chunking, Long.parseLong(length));
   }
 
-  private String getRawUrl(URI baseUrl, String uri) {
+  public URI getRawUri(HttpRequest req) {
+    URI baseUrl = config.baseUri();
+    String uri = req.getUri();
     String rawUrl;
+
     if (uri.startsWith("ws://") || uri.startsWith("wss://") ||
-      uri.startsWith("http://") || uri.startsWith("https://")) {
+        uri.startsWith("http://") || uri.startsWith("https://")) {
       rawUrl = uri;
     } else {
       rawUrl = baseUrl.toString().replaceAll("/$", "") + uri;
     }
 
-    return rawUrl;
-  }
-
-  public URI getRawUri(HttpRequest req) {
-    String rawUrl = getRawUrl(config.baseUri(), req.getUri());
     return URI.create(rawUrl);
   }