Fix spring-webflux header parsing for 7+ (#15502)

Author: jaydeluca

instrumentation/spring/spring-webflux/spring-webflux-5.3/library/src/main/java/io/opentelemetry/instrumentation/spring/webflux/v5_3/WebfluxTextMapGetter.java

@@ -18,7 +18,7 @@ enum WebfluxTextMapGetter implements TextMapGetter<ServerWebExchange> {
 
   @Override
   public Iterable<String> keys(ServerWebExchange exchange) {
-    return exchange.getRequest().getHeaders().keySet();
+    return HeaderUtil.getKeys(exchange.getRequest().getHeaders());
   }
 
   @Nullable

instrumentation/spring/spring-webflux/spring-webflux-5.3/library/src/main/java/io/opentelemetry/instrumentation/spring/webflux/v5_3/internal/HeaderUtil.java

@@ -6,11 +6,14 @@
 package io.opentelemetry.instrumentation.spring.webflux.v5_3.internal;
 
 import static java.util.Collections.emptyList;
+import static java.util.Collections.emptySet;
 
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.MethodType;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
 import java.util.function.Supplier;
 import javax.annotation.Nullable;
 import org.springframework.http.HttpHeaders;
@@ -23,6 +26,7 @@
 public final class HeaderUtil {
 
   @Nullable private static final MethodHandle GET_HEADERS;
+  @Nullable private static final MethodHandle HEADER_NAMES;
 
   static {
     GET_HEADERS =
@@ -32,6 +36,17 @@ public final class HeaderUtil {
             () ->
                 findGetHeadersMethod(
                     MethodType.methodType(List.class, Object.class))); // before spring web 7.0
+
+    // Spring Web 7+
+    MethodHandle headerNames = null;
+    try {
+      headerNames =
+          MethodHandles.lookup()
+              .findVirtual(HttpHeaders.class, "headerNames", MethodType.methodType(Set.class));
+    } catch (Throwable t) {
+      // ignore - will fall back to casting to Map
+    }
+    HEADER_NAMES = headerNames;
   }
 
   @Nullable
@@ -64,5 +79,24 @@ public static List<String> getHeader(HttpHeaders headers, String name) {
     return emptyList();
   }
 
+  @SuppressWarnings("unchecked") // HttpHeaders is a Map in Spring Web 6 and earlier
+  public static Set<String> getKeys(HttpHeaders headers) {
+    if (HEADER_NAMES != null) {
+      // Spring Web 7: HttpHeaders has headerNames() method
+      try {
+        Set<String> result = (Set<String>) HEADER_NAMES.invoke(headers);
+        if (result != null) {
+          return result;
+        }
+      } catch (Throwable t) {
+        // ignore
+      }
+    } else {
+      // Spring Web 6 and earlier: HttpHeaders extends Map
+      return ((Map<String, List<String>>) headers).keySet();
+    }
+    return emptySet();
+  }
+
   private HeaderUtil() {}
 }

instrumentation/spring/spring-webflux/spring-webflux-5.3/library/src/test/java/io/opentelemetry/instrumentation/spring/webflux/v5_3/WebfluxTextMapGetterTest.java

@@ -0,0 +1,90 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.spring.webflux.v5_3;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.ArrayList;
+import java.util.List;
+import org.junit.jupiter.api.Test;
+import org.springframework.mock.http.server.reactive.MockServerHttpRequest;
+import org.springframework.mock.web.server.MockServerWebExchange;
+
+class WebfluxTextMapGetterTest {
+
+  @Test
+  void testGet() {
+    MockServerHttpRequest request =
+        MockServerHttpRequest.get("/test")
+            .header("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01")
+            .header("custom-header", "custom-value")
+            .build();
+
+    MockServerWebExchange exchange = MockServerWebExchange.from(request);
+
+    String traceparent = WebfluxTextMapGetter.INSTANCE.get(exchange, "traceparent");
+    assertThat(traceparent).isEqualTo("00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01");
+
+    String customHeader = WebfluxTextMapGetter.INSTANCE.get(exchange, "custom-header");
+    assertThat(customHeader).isEqualTo("custom-value");
+  }
+
+  @Test
+  void testGetAll() {
+    MockServerHttpRequest request =
+        MockServerHttpRequest.get("/test")
+            .header("accept", "application/json")
+            .header("accept", "text/html")
+            .build();
+
+    MockServerWebExchange exchange = MockServerWebExchange.from(request);
+
+    List<String> acceptHeaders = new ArrayList<>();
+    WebfluxTextMapGetter.INSTANCE.getAll(exchange, "accept").forEachRemaining(acceptHeaders::add);
+
+    assertThat(acceptHeaders).containsExactly("application/json", "text/html");
+  }
+
+  @Test
+  void testKeysWithBaggageHeader() {
+    MockServerHttpRequest request =
+        MockServerHttpRequest.get("/test")
+            .header("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01")
+            .header("baggage", "test-baggage-key-1=test-baggage-value-1")
+            .build();
+
+    MockServerWebExchange exchange = MockServerWebExchange.from(request);
+
+    Iterable<String> keys = WebfluxTextMapGetter.INSTANCE.keys(exchange);
+    assertThat(keys).contains("traceparent", "baggage");
+
+    String baggageValue = WebfluxTextMapGetter.INSTANCE.get(exchange, "baggage");
+    assertThat(baggageValue).isEqualTo("test-baggage-key-1=test-baggage-value-1");
+  }
+
+  @Test
+  void testKeysWithMultipleBaggageHeaders() {
+    MockServerHttpRequest request =
+        MockServerHttpRequest.get("/test")
+            .header("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01")
+            .header("baggage", "test-baggage-key-1=test-baggage-value-1")
+            .header("baggage", "test-baggage-key-2=test-baggage-value-2")
+            .header("x-custom", "custom-value")
+            .build();
+
+    MockServerWebExchange exchange = MockServerWebExchange.from(request);
+
+    Iterable<String> keys = WebfluxTextMapGetter.INSTANCE.keys(exchange);
+    assertThat(keys).contains("traceparent", "baggage", "x-custom");
+
+    List<String> baggageValues = new ArrayList<>();
+    WebfluxTextMapGetter.INSTANCE.getAll(exchange, "baggage").forEachRemaining(baggageValues::add);
+
+    assertThat(baggageValues)
+        .containsExactly(
+            "test-baggage-key-1=test-baggage-value-1", "test-baggage-key-2=test-baggage-value-2");
+  }
+}