Fix headers keySet in WebFlux adapters

Prior to this commit, WebFlux native headers adapters would delegate the
`httpHeaders.keySet` to underlying implementations that do not honor the
`remove*` methods.

This commit fixes the `Set` implementation backing the
`httpHeaders.keySet` and ensures that headers can be safely removed from
the set.

Fixes gh-26361

spring-web/src/main/java/org/springframework/http/client/reactive/JettyHeadersAdapter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 the original author or authors.
+ * Copyright 2002-2021 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -154,7 +154,7 @@ class JettyHeadersAdapter implements MultiValueMap<String, String> {
 
 	@Override
 	public Set<String> keySet() {
-		return this.headers.getFieldNamesCollection();
+		return new HeaderNames();
 	}
 
 	@Override
@@ -227,4 +227,52 @@ class JettyHeadersAdapter implements MultiValueMap<String, String> {
 		}
 	}
 
+
+	private class HeaderNames extends AbstractSet<String> {
+
+		@Override
+		public Iterator<String> iterator() {
+			return new HeaderNamesIterator(headers.getFieldNamesCollection().iterator());
+		}
+
+		@Override
+		public int size() {
+			return headers.getFieldNamesCollection().size();
+		}
+	}
+
+	private final class HeaderNamesIterator implements Iterator<String> {
+
+		private final Iterator<String> iterator;
+
+		@Nullable
+		private String currentName;
+
+		private HeaderNamesIterator(Iterator<String> iterator) {
+			this.iterator = iterator;
+		}
+
+		@Override
+		public boolean hasNext() {
+			return this.iterator.hasNext();
+		}
+
+		@Override
+		public String next() {
+			this.currentName = this.iterator.next();
+			return this.currentName;
+		}
+
+		@Override
+		public void remove() {
+			if (this.currentName == null) {
+				throw new IllegalStateException("No current Header in iterator");
+			}
+			if (!headers.containsKey(this.currentName)) {
+				throw new IllegalStateException("Header not present: " + this.currentName);
+			}
+			headers.remove(this.currentName);
+		}
+	}
+
 }

spring-web/src/main/java/org/springframework/http/client/reactive/NettyHeadersAdapter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 the original author or authors.
+ * Copyright 2002-2021 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -157,7 +157,7 @@ class NettyHeadersAdapter implements MultiValueMap<String, String> {
 
 	@Override
 	public Set<String> keySet() {
-		return this.headers.names();
+		return new HeaderNames();
 	}
 
 	@Override
@@ -230,4 +230,52 @@ class NettyHeadersAdapter implements MultiValueMap<String, String> {
 		}
 	}
 
+
+	private class HeaderNames extends AbstractSet<String> {
+
+		@Override
+		public Iterator<String> iterator() {
+			return new HeaderNamesIterator(headers.names().iterator());
+		}
+
+		@Override
+		public int size() {
+			return headers.names().size();
+		}
+	}
+
+	private final class HeaderNamesIterator implements Iterator<String> {
+
+		private final Iterator<String> iterator;
+
+		@Nullable
+		private String currentName;
+
+		private HeaderNamesIterator(Iterator<String> iterator) {
+			this.iterator = iterator;
+		}
+
+		@Override
+		public boolean hasNext() {
+			return this.iterator.hasNext();
+		}
+
+		@Override
+		public String next() {
+			this.currentName = this.iterator.next();
+			return this.currentName;
+		}
+
+		@Override
+		public void remove() {
+			if (this.currentName == null) {
+				throw new IllegalStateException("No current Header in iterator");
+			}
+			if (!headers.contains(this.currentName)) {
+				throw new IllegalStateException("Header not present: " + this.currentName);
+			}
+			headers.remove(this.currentName);
+		}
+	}
+
 }

spring-web/src/main/java/org/springframework/http/server/reactive/JettyHeadersAdapter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 the original author or authors.
+ * Copyright 2002-2021 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -154,7 +154,7 @@ class JettyHeadersAdapter implements MultiValueMap<String, String> {
 
 	@Override
 	public Set<String> keySet() {
-		return this.headers.getFieldNamesCollection();
+		return new HeaderNames();
 	}
 
 	@Override
@@ -227,4 +227,52 @@ class JettyHeadersAdapter implements MultiValueMap<String, String> {
 		}
 	}
 
+
+	private class HeaderNames extends AbstractSet<String> {
+
+		@Override
+		public Iterator<String> iterator() {
+			return new HeaderNamesIterator(headers.getFieldNamesCollection().iterator());
+		}
+
+		@Override
+		public int size() {
+			return headers.getFieldNamesCollection().size();
+		}
+	}
+
+	private final class HeaderNamesIterator implements Iterator<String> {
+
+		private final Iterator<String> iterator;
+
+		@Nullable
+		private String currentName;
+
+		private HeaderNamesIterator(Iterator<String> iterator) {
+			this.iterator = iterator;
+		}
+
+		@Override
+		public boolean hasNext() {
+			return this.iterator.hasNext();
+		}
+
+		@Override
+		public String next() {
+			this.currentName = this.iterator.next();
+			return this.currentName;
+		}
+
+		@Override
+		public void remove() {
+			if (this.currentName == null) {
+				throw new IllegalStateException("No current Header in iterator");
+			}
+			if (!headers.containsKey(this.currentName)) {
+				throw new IllegalStateException("Header not present: " + this.currentName);
+			}
+			headers.remove(this.currentName);
+		}
+	}
+
 }

spring-web/src/main/java/org/springframework/http/server/reactive/NettyHeadersAdapter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 the original author or authors.
+ * Copyright 2002-2021 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -157,7 +157,7 @@ class NettyHeadersAdapter implements MultiValueMap<String, String> {
 
 	@Override
 	public Set<String> keySet() {
-		return this.headers.names();
+		return new HeaderNames();
 	}
 
 	@Override
@@ -230,4 +230,51 @@ class NettyHeadersAdapter implements MultiValueMap<String, String> {
 		}
 	}
 
+	private class HeaderNames extends AbstractSet<String> {
+
+		@Override
+		public Iterator<String> iterator() {
+			return new HeaderNamesIterator(headers.names().iterator());
+		}
+
+		@Override
+		public int size() {
+			return headers.names().size();
+		}
+	}
+
+	private final class HeaderNamesIterator implements Iterator<String> {
+
+		private final Iterator<String> iterator;
+
+		@Nullable
+		private String currentName;
+
+		private HeaderNamesIterator(Iterator<String> iterator) {
+			this.iterator = iterator;
+		}
+
+		@Override
+		public boolean hasNext() {
+			return this.iterator.hasNext();
+		}
+
+		@Override
+		public String next() {
+			this.currentName = this.iterator.next();
+			return this.currentName;
+		}
+
+		@Override
+		public void remove() {
+			if (this.currentName == null) {
+				throw new IllegalStateException("No current Header in iterator");
+			}
+			if (!headers.contains(this.currentName)) {
+				throw new IllegalStateException("Header not present: " + this.currentName);
+			}
+			headers.remove(this.currentName);
+		}
+	}
+
 }

spring-web/src/main/java/org/springframework/http/server/reactive/TomcatHeadersAdapter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 the original author or authors.
+ * Copyright 2002-2021 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -20,7 +20,6 @@ import java.util.AbstractSet;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Enumeration;
-import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -167,12 +166,7 @@ class TomcatHeadersAdapter implements MultiValueMap<String, String> {
 
 	@Override
 	public Set<String> keySet() {
-		Set<String> result = new HashSet<>(8);
+		return new HeaderNames();
-		Enumeration<String> names = this.headers.names();
-		while (names.hasMoreElements()) {
-			result.add(names.nextElement());
-		}
-		return result;
 	}
 
 	@Override
@@ -247,4 +241,59 @@ class TomcatHeadersAdapter implements MultiValueMap<String, String> {
 		}
 	}
 
+
+	private class HeaderNames extends AbstractSet<String> {
+
+		@Override
+		public Iterator<String> iterator() {
+			return new HeaderNamesIterator(headers.names());
+		}
+
+		@Override
+		public int size() {
+			Enumeration<String> names = headers.names();
+			int size = 0;
+			while (names.hasMoreElements()) {
+				names.nextElement();
+				size++;
+			}
+			return size;
+		}
+	}
+
+	private final class HeaderNamesIterator implements Iterator<String> {
+
+		private final Enumeration<String> enumeration;
+
+		@Nullable
+		private String currentName;
+
+		private HeaderNamesIterator(Enumeration<String> enumeration) {
+			this.enumeration = enumeration;
+		}
+
+		@Override
+		public boolean hasNext() {
+			return this.enumeration.hasMoreElements();
+		}
+
+		@Override
+		public String next() {
+			this.currentName = this.enumeration.nextElement();
+			return this.currentName;
+		}
+
+		@Override
+		public void remove() {
+			if (this.currentName == null) {
+				throw new IllegalStateException("No current Header in iterator");
+			}
+			int index = headers.findHeader(this.currentName, 0);
+			if (index == -1) {
+				throw new IllegalStateException("Header not present: " + this.currentName);
+			}
+			headers.removeHeader(index);
+		}
+	}
+
 }

spring-web/src/main/java/org/springframework/http/server/reactive/UndertowHeadersAdapter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 the original author or authors.
+ * Copyright 2002-2021 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -149,9 +149,7 @@ class UndertowHeadersAdapter implements MultiValueMap<String, String> {
 
 	@Override
 	public Set<String> keySet() {
-		return this.headers.getHeaderNames().stream()
+		return new HeaderNames();
-				.map(HttpString::toString)
-				.collect(Collectors.toSet());
 	}
 
 	@Override
@@ -225,4 +223,52 @@ class UndertowHeadersAdapter implements MultiValueMap<String, String> {
 		}
 	}
 
+
+	private class HeaderNames extends AbstractSet<String> {
+
+		@Override
+		public Iterator<String> iterator() {
+			return new HeaderNamesIterator(headers.getHeaderNames().iterator());
+		}
+
+		@Override
+		public int size() {
+			return headers.getHeaderNames().size();
+		}
+	}
+
+	private final class HeaderNamesIterator implements Iterator<String> {
+
+		private final Iterator<HttpString> iterator;
+
+		@Nullable
+		private String currentName;
+
+		private HeaderNamesIterator(Iterator<HttpString> iterator) {
+			this.iterator = iterator;
+		}
+
+		@Override
+		public boolean hasNext() {
+			return this.iterator.hasNext();
+		}
+
+		@Override
+		public String next() {
+			this.currentName = this.iterator.next().toString();
+			return this.currentName;
+		}
+
+		@Override
+		public void remove() {
+			if (this.currentName == null) {
+				throw new IllegalStateException("No current Header in iterator");
+			}
+			if (!headers.contains(this.currentName)) {
+				throw new IllegalStateException("Header not present: " + this.currentName);
+			}
+			headers.remove(this.currentName);
+		}
+	}
+
 }

spring-web/src/test/java/org/springframework/http/server/reactive/HeadersAdaptersTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2019 the original author or authors.
+ * Copyright 2002-2021 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -21,6 +21,7 @@ import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 import java.util.Arrays;
+import java.util.Iterator;
 import java.util.Locale;
 import java.util.stream.Stream;
 
@@ -37,6 +38,7 @@ import org.springframework.util.LinkedCaseInsensitiveMap;
 import org.springframework.util.MultiValueMap;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.junit.jupiter.params.provider.Arguments.arguments;
 
 /**
@@ -102,6 +104,25 @@ class HeadersAdaptersTests {
 		assertThat(headers.getFirst("TestHeader")).isNull();
 	}
 
+	@ParameterizedHeadersTest
+	void shouldReflectChangesOnKeyset(String displayName, MultiValueMap<String, String> headers) {
+		headers.add("TestHeader", "first");
+		assertThat(headers.keySet()).hasSize(1);
+		headers.keySet().removeIf("TestHeader"::equals);
+		assertThat(headers.keySet()).hasSize(0);
+	}
+
+	@ParameterizedHeadersTest
+	void shouldFailIfHeaderRemovedFromKeyset(String displayName, MultiValueMap<String, String> headers) {
+		headers.add("TestHeader", "first");
+		assertThat(headers.keySet()).hasSize(1);
+		Iterator<String> names = headers.keySet().iterator();
+		assertThat(names.hasNext()).isTrue();
+		assertThat(names.next()).isEqualTo("TestHeader");
+		names.remove();
+		assertThatThrownBy(names::remove).isInstanceOf(IllegalStateException.class);
+	}
+
 	@Retention(RetentionPolicy.RUNTIME)
 	@Target(ElementType.METHOD)
 	@ParameterizedTest(name = "[{index}] {0}")