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
This commit is contained in:
Brian Clozel 2021-01-11 09:18:10 +01:00
parent 689b5566bf
commit e4dc863ad0
7 changed files with 328 additions and 21 deletions

View File

@ -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);
}
}
}

View File

@ -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);
}
}
}

View File

@ -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);
}
}
}

View File

@ -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);
}
}
}

View File

@ -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);
Enumeration<String> names = this.headers.names();
while (names.hasMoreElements()) {
result.add(names.nextElement());
}
return result;
return new HeaderNames();
}
@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);
}
}
}

View File

@ -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()
.map(HttpString::toString)
.collect(Collectors.toSet());
return new HeaderNames();
}
@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);
}
}
}

View File

@ -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}")