[java] ensure all headers are removed or replaced

Author: joerg1985

java/src/org/openqa/selenium/grid/node/DefaultActiveSession.java

@@ -5,8 +5,6 @@
 import java.io.UncheckedIOException;
 import java.net.URL;
 import java.time.Instant;
-import java.util.stream.Collectors;
-import java.util.stream.StreamSupport;
 import org.openqa.selenium.Capabilities;
 import org.openqa.selenium.grid.web.ReverseProxyHandler;
 import org.openqa.selenium.internal.Require;
@@ -43,12 +41,7 @@ protected DefaultActiveSession(
 
   @Override
   public HttpResponse execute(HttpRequest req) throws UncheckedIOException {
-    String host = "host";
-    StreamSupport.stream(req.getHeaderNames().spliterator(), true)
-        .filter(host::equalsIgnoreCase)
-        .collect(Collectors.toList())
-        .forEach(req::removeHeader);
-    req.addHeader(host, String.format("%s:%s", getUri().getHost(), getUri().getPort()));
+    req.setHeader("host", String.format("%s:%s", getUri().getHost(), getUri().getPort()));
     HttpResponse res = handler.execute(req);
     if (req.getMethod() == DELETE && killUrl.equals(req.getUri())) {
       stop();

java/src/org/openqa/selenium/remote/http/HttpMessage.java

@@ -26,8 +26,8 @@
 import com.google.common.net.MediaType;
 import java.io.InputStream;
 import java.nio.charset.Charset;
+import java.util.Collections;
 import java.util.HashMap;
-import java.util.Iterator;
 import java.util.Map;
 import java.util.Objects;
 import java.util.function.Supplier;
@@ -66,7 +66,7 @@ public Iterable<String> getAttributeNames() {
   }
 
   public Iterable<String> getHeaderNames() {
-    return headers.keySet();
+    return Collections.unmodifiableCollection(headers.keySet());
   }
 
   public Iterable<String> getHeaders(String name) {
@@ -78,16 +78,12 @@ public Iterable<String> getHeaders(String name) {
   }
 
   public String getHeader(String name) {
-    Iterable<String> initialHeaders = getHeaders(name);
-    if (initialHeaders == null) {
-      return null;
-    }
-
-    Iterator<String> headers = initialHeaders.iterator();
-    if (headers.hasNext()) {
-      return headers.next();
-    }
-    return null;
+    return headers.entries().stream()
+        .filter(e -> Objects.nonNull(e.getKey()))
+        .filter(e -> e.getKey().equalsIgnoreCase(name.toLowerCase()))
+        .map(Map.Entry::getValue)
+        .findFirst()
+        .orElse(null);
   }
 
   public M setHeader(String name, String value) {
@@ -100,12 +96,7 @@ public M addHeader(String name, String value) {
   }
 
   public M removeHeader(String name) {
-    String toRemove =
-        headers.keySet().stream()
-            .filter(header -> header.equalsIgnoreCase(name))
-            .findFirst()
-            .orElse(name);
-    headers.removeAll(toRemove);
+    headers.keySet().removeIf(header -> header.equalsIgnoreCase(name));
     return self();
   }
 

java/test/org/openqa/selenium/remote/http/HttpMessageTest.java

@@ -0,0 +1,96 @@
+// 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.http;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+
+@Tag("UnitTests")
+class HttpMessageTest {
+
+  @Test
+  void allHeadersAreAdded() {
+    for (HttpMessage<?> message :
+        Arrays.asList(new HttpRequest(HttpMethod.GET, "/"), new HttpResponse())) {
+      message.addHeader("Content-Length", "1024");
+      message.addHeader("Content-length", "2048");
+      message.addHeader("content-length", "4096");
+
+      Set<String> headers = new HashSet<>();
+
+      message.getHeaderNames().forEach(headers::add);
+
+      assertThat(headers.contains("Content-Length")).isTrue();
+      assertThat(headers.contains("Content-length")).isTrue();
+      assertThat(headers.contains("content-length")).isTrue();
+    }
+  }
+
+  @Test
+  void readingIsCaseInsensitive() {
+    for (HttpMessage<?> message :
+        Arrays.asList(new HttpRequest(HttpMethod.GET, "/"), new HttpResponse())) {
+      message.addHeader("Content-Length", "1024");
+      message.addHeader("Content-length", "2048");
+      message.addHeader("content-length", "4096");
+
+      assertThat(message.getHeader("Content-Length")).isEqualTo("4096");
+    }
+  }
+
+  @Test
+  void replacingIsCaseInsensitive() {
+    for (HttpMessage<?> message :
+        Arrays.asList(new HttpRequest(HttpMethod.GET, "/"), new HttpResponse())) {
+      message.addHeader("Content-Length", "1024");
+      message.addHeader("Content-length", "2048");
+      message.addHeader("content-length", "4096");
+      message.setHeader("contenT-length", "8192");
+
+      Set<String> headers = new HashSet<>();
+
+      message.getHeaderNames().forEach(headers::add);
+
+      assertThat(message.getHeader("content-length")).isEqualTo("8192");
+      assertThat(headers.contains("Content-Length")).isFalse();
+      assertThat(headers.contains("Content-length")).isFalse();
+      assertThat(headers.contains("content-length")).isFalse();
+    }
+  }
+
+  @Test
+  void allHeadersAreRemoved() {
+    for (HttpMessage<?> message :
+        Arrays.asList(new HttpRequest(HttpMethod.GET, "/"), new HttpResponse())) {
+      message.addHeader("Content-Length", "1024");
+      message.addHeader("Content-length", "2048");
+      message.addHeader("content-length", "4096");
+
+      assertThat(message.getHeaderNames().iterator().hasNext()).isTrue();
+
+      message.removeHeader("Content-Length");
+
+      assertThat(message.getHeaderNames().iterator().hasNext()).isFalse();
+    }
+  }
+}