From 2b328c62008cada338451979c0a6afa89d34dd04 Mon Sep 17 00:00:00 2001 From: Haotian Zhang <928016560@qq.com> Date: Thu, 9 Jul 2026 19:10:47 +0800 Subject: [PATCH] fix: inject baggage as W3C header for reactive clients to avoid OTel scope leak Co-Authored-By: Claude Sonnet 4.6 (200K context) Signed-off-by: Haotian Zhang <928016560@qq.com> --- .../trace/TraceClientPreEnhancedPlugin.java | 36 +++ .../scg/EnhancedGatewayGlobalFilter.java | 47 +++- ...hancedWebClientExchangeFilterFunction.java | 27 ++- .../plugin/EnhancedContextKeys.java | 41 ++++ .../util/OtelBaggageScopeHelper.java | 208 ++++++++++++++++++ .../util/OtelBaggageScopeHelperTest.java | 143 ++++++++++++ 6 files changed, 494 insertions(+), 8 deletions(-) create mode 100644 spring-cloud-tencent-rpc-enhancement/src/main/java/com/tencent/cloud/rpc/enhancement/plugin/EnhancedContextKeys.java create mode 100644 spring-cloud-tencent-rpc-enhancement/src/main/java/com/tencent/cloud/rpc/enhancement/util/OtelBaggageScopeHelper.java create mode 100644 spring-cloud-tencent-rpc-enhancement/src/test/java/com/tencent/cloud/rpc/enhancement/util/OtelBaggageScopeHelperTest.java diff --git a/spring-cloud-tencent-plugin-starters/spring-cloud-starter-tencent-trace-plugin/src/main/java/com/tencent/cloud/plugin/trace/TraceClientPreEnhancedPlugin.java b/spring-cloud-tencent-plugin-starters/spring-cloud-starter-tencent-trace-plugin/src/main/java/com/tencent/cloud/plugin/trace/TraceClientPreEnhancedPlugin.java index 52cc7e348..0a3fb6f88 100644 --- a/spring-cloud-tencent-plugin-starters/spring-cloud-starter-tencent-trace-plugin/src/main/java/com/tencent/cloud/plugin/trace/TraceClientPreEnhancedPlugin.java +++ b/spring-cloud-tencent-plugin-starters/spring-cloud-starter-tencent-trace-plugin/src/main/java/com/tencent/cloud/plugin/trace/TraceClientPreEnhancedPlugin.java @@ -23,6 +23,7 @@ import java.util.Map; import com.tencent.cloud.plugin.trace.attribute.SpanAttributesProvider; import com.tencent.cloud.polaris.context.PolarisSDKContextManager; +import com.tencent.cloud.rpc.enhancement.plugin.EnhancedContextKeys; import com.tencent.cloud.rpc.enhancement.plugin.EnhancedPlugin; import com.tencent.cloud.rpc.enhancement.plugin.EnhancedPluginContext; import com.tencent.cloud.rpc.enhancement.plugin.EnhancedPluginType; @@ -58,7 +59,21 @@ public class TraceClientPreEnhancedPlugin implements EnhancedPlugin { } } } + if (CollectionUtils.isEmpty(attributes)) { + return; + } + + // Reactive scenarios (SCG / WebClient): only stage the baggage attributes to be written; the + // downstream interceptor injects them as a W3C baggage header on the outgoing request, so no + // OTel Scope is attached here and there is no ThreadLocal lifecycle to close across async + // boundaries. + if (isReactiveScenario(context)) { + context.getExtraData().put(EnhancedContextKeys.PENDING_BAGGAGE_ATTRIBUTES_KEY, attributes); + return; + } + // Blocking scenarios (Feign / RestTemplate): keep the original semantics, attaching to the + // current thread so the FINALLY stage can close it. TraceAttributes traceAttributes = new TraceAttributes(); traceAttributes.setAttributes(attributes); traceAttributes.setAttributeLocation(TraceAttributes.AttributeLocation.BAGGAGE); @@ -71,6 +86,27 @@ public class TraceClientPreEnhancedPlugin implements EnhancedPlugin { } } + /** + * Whether current invocation is a reactive client scenario (SCG / WebClient). Matched by class + * name prefix to avoid a hard dependency on spring-webflux / spring-cloud-gateway. + * + * @param context enhanced plugin context + * @return true if the current call is a reactive scenario + */ + private boolean isReactiveScenario(EnhancedPluginContext context) { + Object origin = context.getOriginRequest(); + if (origin == null) { + return false; + } + String className = origin.getClass().getName(); + // SCG: org.springframework.web.server.ServerWebExchange implementation + // WebClient: org.springframework.web.reactive.function.client.ClientRequest implementation + return className.startsWith("org.springframework.web.server.") + || className.startsWith("org.springframework.mock.web.server.") + || className.startsWith("org.springframework.web.reactive.function.client.") + || className.startsWith("org.springframework.cloud.gateway."); + } + @Override public int getOrder() { return PluginOrderConstant.ClientPluginOrder.TRACE_CLIENT_PLUGIN_ORDER; diff --git a/spring-cloud-tencent-rpc-enhancement/src/main/java/com/tencent/cloud/rpc/enhancement/instrument/scg/EnhancedGatewayGlobalFilter.java b/spring-cloud-tencent-rpc-enhancement/src/main/java/com/tencent/cloud/rpc/enhancement/instrument/scg/EnhancedGatewayGlobalFilter.java index 721e8fd64..003916a0e 100644 --- a/spring-cloud-tencent-rpc-enhancement/src/main/java/com/tencent/cloud/rpc/enhancement/instrument/scg/EnhancedGatewayGlobalFilter.java +++ b/spring-cloud-tencent-rpc-enhancement/src/main/java/com/tencent/cloud/rpc/enhancement/instrument/scg/EnhancedGatewayGlobalFilter.java @@ -31,6 +31,7 @@ import com.tencent.cloud.rpc.enhancement.plugin.EnhancedPluginType; import com.tencent.cloud.rpc.enhancement.plugin.EnhancedRequestContext; import com.tencent.cloud.rpc.enhancement.plugin.EnhancedResponseContext; import com.tencent.cloud.rpc.enhancement.util.EnhancedPluginUtils; +import com.tencent.cloud.rpc.enhancement.util.OtelBaggageScopeHelper; import com.tencent.polaris.api.utils.CollectionUtils; import com.tencent.polaris.api.utils.StringUtils; import com.tencent.polaris.circuitbreak.client.exception.CallAbortedException; @@ -163,11 +164,23 @@ public class EnhancedGatewayGlobalFilter implements GlobalFilter, Ordered { } // Exchange may be changed in plugin ServerWebExchange exchange = (ServerWebExchange) enhancedPluginContext.getOriginRequest(); - return chain.filter(exchange) - .doOnSubscribe(v -> { - setTargetServiceInstance(exchange, enhancedPluginContext, serviceId); - pluginRunner.run(EnhancedPluginType.Client.BEFORE_CALLING, enhancedPluginContext); - }) + + // Run BEFORE_CALLING synchronously, before chain.filter, instead of inside doOnSubscribe as + // before. TraceClientPreEnhancedPlugin stages the baggage attributes into + // enhancedPluginContext.extraData (PENDING_BAGGAGE_ATTRIBUTES_KEY) here so they can be written + // into the downstream request headers below, which must happen before the request is sent. + // Moving it earlier is safe today: setTargetServiceInstance only needs GATEWAY_REQUEST_URL_ATTR, + // already rewritten by ReactiveLoadBalancerClientFilter (lower order, runs first), and the only + // reactive BEFORE_CALLING plugin is the trace plugin. A future plugin that relies on + // subscribe-time context would observe a different timing. + setTargetServiceInstance(exchange, enhancedPluginContext, serviceId); + pluginRunner.run(EnhancedPluginType.Client.BEFORE_CALLING, enhancedPluginContext); + + // Inject the pending baggage as a W3C baggage header on the downstream request before + // subscribing. This avoids any dependency on the OTel Context ThreadLocal lifecycle and + // prevents Scope.close from failing across async boundaries. + ServerWebExchange downstreamExchange = injectPendingBaggageHeader(exchange, enhancedPluginContext); + return chain.filter(downstreamExchange) .doOnSuccess(v -> { MetadataContext metadataContextOnSuccess = originExchange.getAttribute( MetadataConstant.HeaderName.METADATA_CONTEXT); @@ -177,8 +190,8 @@ public class EnhancedGatewayGlobalFilter implements GlobalFilter, Ordered { enhancedPluginContext.setDelay(System.currentTimeMillis() - startTime); EnhancedResponseContext enhancedResponseContext = EnhancedResponseContext.builder() - .httpStatus(exchange.getResponse().getRawStatusCode()) - .httpHeaders(exchange.getResponse().getHeaders()) + .httpStatus(downstreamExchange.getResponse().getRawStatusCode()) + .httpHeaders(downstreamExchange.getResponse().getHeaders()) .build(); enhancedPluginContext.setResponse(enhancedResponseContext); @@ -210,6 +223,26 @@ public class EnhancedGatewayGlobalFilter implements GlobalFilter, Ordered { }); } + /** + * Inject pending baggage attributes into the downstream request as a W3C baggage HTTP header. + * This avoids any dependency on the OTel Context ThreadLocal lifecycle and prevents Scope.close + * from failing across async boundaries. + * + * @param exchange original ServerWebExchange + * @param ctx enhanced plugin context (holds the baggage attributes to write) + * @return the original exchange when there is nothing to write, otherwise a mutated exchange + */ + private ServerWebExchange injectPendingBaggageHeader(ServerWebExchange exchange, EnhancedPluginContext ctx) { + String existing = exchange.getRequest().getHeaders().getFirst(OtelBaggageScopeHelper.BAGGAGE_HEADER); + String merged = OtelBaggageScopeHelper.resolvePendingBaggageHeader(ctx, existing); + if (merged == null) { + return exchange; + } + return exchange.mutate().request(builder -> builder.headers(headers -> + headers.set(OtelBaggageScopeHelper.BAGGAGE_HEADER, merged) + )).build(); + } + @Override public int getOrder() { return OrderConstant.Client.Scg.ENHANCED_FILTER_ORDER; diff --git a/spring-cloud-tencent-rpc-enhancement/src/main/java/com/tencent/cloud/rpc/enhancement/instrument/webclient/EnhancedWebClientExchangeFilterFunction.java b/spring-cloud-tencent-rpc-enhancement/src/main/java/com/tencent/cloud/rpc/enhancement/instrument/webclient/EnhancedWebClientExchangeFilterFunction.java index 1cfc6bf8e..39716beaf 100644 --- a/spring-cloud-tencent-rpc-enhancement/src/main/java/com/tencent/cloud/rpc/enhancement/instrument/webclient/EnhancedWebClientExchangeFilterFunction.java +++ b/spring-cloud-tencent-rpc-enhancement/src/main/java/com/tencent/cloud/rpc/enhancement/instrument/webclient/EnhancedWebClientExchangeFilterFunction.java @@ -30,6 +30,7 @@ import com.tencent.cloud.rpc.enhancement.plugin.EnhancedPluginType; import com.tencent.cloud.rpc.enhancement.plugin.EnhancedRequestContext; import com.tencent.cloud.rpc.enhancement.plugin.EnhancedResponseContext; import com.tencent.cloud.rpc.enhancement.util.EnhancedPluginUtils; +import com.tencent.cloud.rpc.enhancement.util.OtelBaggageScopeHelper; import com.tencent.polaris.api.utils.CollectionUtils; import com.tencent.polaris.circuitbreak.client.exception.CallAbortedException; import com.tencent.polaris.fault.client.exception.FaultInjectionException; @@ -134,7 +135,11 @@ public class EnhancedWebClientExchangeFilterFunction implements ExchangeFilterFu } // request may be changed by plugin ClientRequest request = (ClientRequest) enhancedPluginContext.getOriginRequest(); - return next.exchange(request) + // Inject the baggage attributes staged in the pre stage as a W3C baggage header on the + // downstream request. This avoids any dependency on the OTel Context ThreadLocal lifecycle + // and prevents Scope.close from failing across async boundaries. + ClientRequest downstreamRequest = injectPendingBaggageHeader(request, enhancedPluginContext); + return next.exchange(downstreamRequest) .doOnSuccess(response -> { enhancedPluginContext.setDelay(System.currentTimeMillis() - startTime); @@ -160,6 +165,26 @@ public class EnhancedWebClientExchangeFilterFunction implements ExchangeFilterFu }); } + /** + * Inject pending baggage attributes into the downstream ClientRequest as a W3C baggage HTTP + * header. This avoids any dependency on the OTel Context ThreadLocal lifecycle and prevents + * Scope.close from failing across async boundaries. + * + * @param request original ClientRequest + * @param ctx enhanced plugin context (holds the baggage attributes to write) + * @return the original request when there is nothing to write, otherwise a mutated ClientRequest + */ + private ClientRequest injectPendingBaggageHeader(ClientRequest request, EnhancedPluginContext ctx) { + String existing = request.headers().getFirst(OtelBaggageScopeHelper.BAGGAGE_HEADER); + String merged = OtelBaggageScopeHelper.resolvePendingBaggageHeader(ctx, existing); + if (merged == null) { + return request; + } + return ClientRequest.from(request) + .headers(headers -> headers.set(OtelBaggageScopeHelper.BAGGAGE_HEADER, merged)) + .build(); + } + private URI getServiceUri(ClientRequest clientRequest) { Object instance = MetadataContextHolder.get() .getLoadbalancerMetadata().get(LOAD_BALANCER_SERVICE_INSTANCE); diff --git a/spring-cloud-tencent-rpc-enhancement/src/main/java/com/tencent/cloud/rpc/enhancement/plugin/EnhancedContextKeys.java b/spring-cloud-tencent-rpc-enhancement/src/main/java/com/tencent/cloud/rpc/enhancement/plugin/EnhancedContextKeys.java new file mode 100644 index 000000000..e37f3a4ff --- /dev/null +++ b/spring-cloud-tencent-rpc-enhancement/src/main/java/com/tencent/cloud/rpc/enhancement/plugin/EnhancedContextKeys.java @@ -0,0 +1,41 @@ +/* + * Tencent is pleased to support the open source community by making spring-cloud-tencent available. + * + * Copyright (C) 2021 Tencent. All rights reserved. + * + * Licensed under the BSD 3-Clause License (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://opensource.org/licenses/BSD-3-Clause + * + * 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 com.tencent.cloud.rpc.enhancement.plugin; + +/** + * Common key constants stored in {@link EnhancedPluginContext#getExtraData()}. + * + *

These constants live in the rpc-enhancement module so they can be referenced by both the + * SCG / WebClient reactive client interceptors and the upper trace-plugin, without creating a + * module cycle.

+ * + * @author Haotian Zhang + */ +public final class EnhancedContextKeys { + + /** + * Pending baggage attributes (Map of string) staged by the trace plugin in the BEFORE_CALLING + * stage for reactive clients (SCG / WebClient). The SCG / WebClient interceptor removes them and + * injects them as a W3C baggage HTTP header on the outgoing request. No OTel Scope is attached, + * so there is no ThreadLocal lifecycle to close across async boundaries. + */ + public static final String PENDING_BAGGAGE_ATTRIBUTES_KEY = "PENDING_BAGGAGE_ATTRIBUTES_KEY"; + + private EnhancedContextKeys() { + } +} diff --git a/spring-cloud-tencent-rpc-enhancement/src/main/java/com/tencent/cloud/rpc/enhancement/util/OtelBaggageScopeHelper.java b/spring-cloud-tencent-rpc-enhancement/src/main/java/com/tencent/cloud/rpc/enhancement/util/OtelBaggageScopeHelper.java new file mode 100644 index 000000000..f1f7994d9 --- /dev/null +++ b/spring-cloud-tencent-rpc-enhancement/src/main/java/com/tencent/cloud/rpc/enhancement/util/OtelBaggageScopeHelper.java @@ -0,0 +1,208 @@ +/* + * Tencent is pleased to support the open source community by making spring-cloud-tencent available. + * + * Copyright (C) 2021 Tencent. All rights reserved. + * + * Licensed under the BSD 3-Clause License (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://opensource.org/licenses/BSD-3-Clause + * + * 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 com.tencent.cloud.rpc.enhancement.util; + +import java.io.ByteArrayOutputStream; +import java.nio.charset.StandardCharsets; +import java.util.BitSet; +import java.util.LinkedHashMap; +import java.util.Map; + +import com.tencent.cloud.rpc.enhancement.plugin.EnhancedContextKeys; +import com.tencent.cloud.rpc.enhancement.plugin.EnhancedPluginContext; + +/** + * Helper that builds a W3C baggage HTTP header value from a map of attributes. + * + *

The previous approach reflectively called {@code Baggage.makeCurrent()} to attach baggage + * to the current thread's OTel Context and closed it after the downstream request was sent. In + * Reactor / Netty async scenarios, attach and close are not guaranteed to run on the same thread, + * and the OTel Agent also restores the ThreadLocal when its own Scope closes. Both cause the + * {@code current() != toAttach} check inside {@code ScopeImpl.close()} to fail, so the scope is + * never actually released.

+ * + *

To avoid this entirely, this helper does not touch {@code ThreadLocalContextStorage} at all. + * Instead it produces a W3C baggage compliant header value. Callers write it into the downstream + * request headers before sending; the OTel Agent on the server side parses it back into + * {@code Baggage}, which is semantically identical to attach but has no ThreadLocal lifecycle.

+ * + * @author Haotian Zhang + */ +public final class OtelBaggageScopeHelper { + + /** + * Standard W3C baggage HTTP header name. + */ + public static final String BAGGAGE_HEADER = "baggage"; + + private static final char[] HEX = "0123456789ABCDEF".toCharArray(); + + /** + * The {@code baggage-octet} set from the W3C Baggage spec: US-ASCII characters excluding CTLs, + * whitespace, DQUOTE, comma, semicolon and backslash. Any code point outside this set must be + * percent-encoded. This intentionally differs from {@code x-www-form-urlencoded}: a space becomes + * {@code %20}, not {@code +}, because OTel's {@code BaggageCodec} only decodes {@code %XX} and + * keeps {@code +} as a literal plus. + */ + private static final BitSet BAGGAGE_OCTET = new BitSet(128); + + static { + for (int c = 0x21; c <= 0x7E; c++) { + BAGGAGE_OCTET.set(c); + } + BAGGAGE_OCTET.clear('"'); + BAGGAGE_OCTET.clear(','); + BAGGAGE_OCTET.clear(';'); + BAGGAGE_OCTET.clear('\\'); + } + + private OtelBaggageScopeHelper() { + } + + /** + * Encode the given attributes as a W3C baggage header value. + * + *

Format is {@code key1=value1,key2=value2}; both key and value are percent-encoded per + * RFC 3986. Returns {@code null} when the input is empty.

+ * + * @param attributes attributes to encode + * @return W3C baggage header value or null when input is empty + */ + public static String encodeBaggage(Map attributes) { + if (attributes == null || attributes.isEmpty()) { + return null; + } + StringBuilder sb = new StringBuilder(); + for (Map.Entry e : attributes.entrySet()) { + if (e.getKey() == null || e.getValue() == null) { + continue; + } + if (sb.length() > 0) { + sb.append(','); + } + sb.append(percentEncode(e.getKey())).append('=').append(percentEncode(e.getValue())); + } + return sb.length() == 0 ? null : sb.toString(); + } + + /** + * Merge the given attributes into an existing baggage header value. + * + *

Existing keys are overwritten by the new value; other keys are preserved. Returns the + * input unchanged when no attributes are provided.

+ * + * @param existing existing baggage header value (nullable) + * @param attributes attributes to merge in + * @return merged baggage header value, or null when both inputs empty + */ + public static String mergeIntoBaggageHeader(String existing, Map attributes) { + if (existing == null || existing.isEmpty()) { + return encodeBaggage(attributes); + } + Map merged = new LinkedHashMap<>(); + for (String pair : existing.split(",")) { + String trimmed = pair.trim(); + if (trimmed.isEmpty()) { + continue; + } + int semi = trimmed.indexOf(';'); + String kv = semi >= 0 ? trimmed.substring(0, semi) : trimmed; + int eq = kv.indexOf('='); + if (eq <= 0) { + continue; + } + merged.put(percentDecode(kv.substring(0, eq).trim()), percentDecode(kv.substring(eq + 1).trim())); + } + if (attributes != null) { + for (Map.Entry e : attributes.entrySet()) { + if (e.getKey() != null && e.getValue() != null) { + merged.put(e.getKey(), e.getValue()); + } + } + } + return encodeBaggage(merged); + } + + /** + * Resolve the pending baggage staged in the context into a baggage header value to write on the + * downstream request. Removes {@link EnhancedContextKeys#PENDING_BAGGAGE_ATTRIBUTES_KEY} from the + * context and merges it with the request's existing baggage header. Returns {@code null} when + * there is nothing to write, so callers can leave the request untouched. + * + * @param ctx enhanced plugin context holding the staged baggage attributes + * @param existingHeader the request's current baggage header value (nullable) + * @return the merged baggage header value, or null when there is nothing to write + */ + @SuppressWarnings("unchecked") + public static String resolvePendingBaggageHeader(EnhancedPluginContext ctx, String existingHeader) { + Map pending = (Map) ctx.getExtraData() + .remove(EnhancedContextKeys.PENDING_BAGGAGE_ATTRIBUTES_KEY); + if (pending == null || pending.isEmpty()) { + return null; + } + return mergeIntoBaggageHeader(existingHeader, pending); + } + + private static String percentEncode(String s) { + // Common case: plain ASCII keys/values (trace ids, hex) are all baggage-octet, so return the + // input without touching the byte[] / StringBuilder path. + if (needsNoEncoding(s)) { + return s; + } + byte[] bytes = s.getBytes(StandardCharsets.UTF_8); + StringBuilder sb = new StringBuilder(bytes.length); + for (byte value : bytes) { + int b = value & 0xFF; + if (b < 128 && BAGGAGE_OCTET.get(b)) { + sb.append((char) b); + } + else { + sb.append('%').append(HEX[b >> 4]).append(HEX[b & 0x0F]); + } + } + return sb.toString(); + } + + private static boolean needsNoEncoding(String s) { + for (int i = 0; i < s.length(); i++) { + char c = s.charAt(i); + if (c >= 128 || !BAGGAGE_OCTET.get(c)) { + return false; + } + } + return true; + } + + private static String percentDecode(String s) { + ByteArrayOutputStream buffer = new ByteArrayOutputStream(s.length()); + for (int i = 0; i < s.length(); i++) { + char c = s.charAt(i); + if (c == '%' && i + 2 < s.length()) { + int high = Character.digit(s.charAt(i + 1), 16); + int low = Character.digit(s.charAt(i + 2), 16); + if (high >= 0 && low >= 0) { + buffer.write((high << 4) + low); + i += 2; + continue; + } + } + buffer.write(c); + } + return new String(buffer.toByteArray(), StandardCharsets.UTF_8); + } +} diff --git a/spring-cloud-tencent-rpc-enhancement/src/test/java/com/tencent/cloud/rpc/enhancement/util/OtelBaggageScopeHelperTest.java b/spring-cloud-tencent-rpc-enhancement/src/test/java/com/tencent/cloud/rpc/enhancement/util/OtelBaggageScopeHelperTest.java new file mode 100644 index 000000000..972bba005 --- /dev/null +++ b/spring-cloud-tencent-rpc-enhancement/src/test/java/com/tencent/cloud/rpc/enhancement/util/OtelBaggageScopeHelperTest.java @@ -0,0 +1,143 @@ +/* + * Tencent is pleased to support the open source community by making spring-cloud-tencent available. + * + * Copyright (C) 2021 Tencent. All rights reserved. + * + * Licensed under the BSD 3-Clause License (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://opensource.org/licenses/BSD-3-Clause + * + * 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 com.tencent.cloud.rpc.enhancement.util; + +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.Map; + +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Unit test for {@link OtelBaggageScopeHelper}. + */ +public class OtelBaggageScopeHelperTest { + + @Test + public void testEncodeBaggage_empty() { + assertThat(OtelBaggageScopeHelper.encodeBaggage(null)).isNull(); + assertThat(OtelBaggageScopeHelper.encodeBaggage(new HashMap<>())).isNull(); + } + + @Test + public void testEncodeBaggage_simple() { + Map attrs = new LinkedHashMap<>(); + attrs.put("k1", "v1"); + attrs.put("k2", "v2"); + + String header = OtelBaggageScopeHelper.encodeBaggage(attrs); + assertThat(header).isEqualTo("k1=v1,k2=v2"); + } + + @Test + public void testEncodeBaggage_specialChars() { + Map attrs = new LinkedHashMap<>(); + attrs.put("custom.k", "v with space"); + attrs.put("namespace-id", "ns-a,b"); + + String header = OtelBaggageScopeHelper.encodeBaggage(attrs); + // space -> %20 (RFC 3986, not x-www-form-urlencoded '+'), comma -> %2C, dot kept as-is + assertThat(header).contains("custom.k=v%20with%20space"); + assertThat(header).contains("namespace-id=ns-a%2Cb"); + } + + @Test + public void testEncodeBaggage_skipNullEntry() { + Map attrs = new LinkedHashMap<>(); + attrs.put("k1", "v1"); + attrs.put("k2", null); + attrs.put(null, "v3"); + attrs.put("k4", "v4"); + + String header = OtelBaggageScopeHelper.encodeBaggage(attrs); + assertThat(header).isEqualTo("k1=v1,k4=v4"); + } + + @Test + public void testEncodeBaggage_spaceUsesPercent20NotPlus() { + Map attrs = new LinkedHashMap<>(); + attrs.put("k", "a b"); + + String header = OtelBaggageScopeHelper.encodeBaggage(attrs); + // OTel BaggageCodec decodes only %XX and keeps '+' literal, so a space must be %20, never '+'. + assertThat(header).isEqualTo("k=a%20b"); + assertThat(header).doesNotContain("+"); + } + + @Test + public void testMergeIntoBaggageHeader_percentEncodedExistingRoundTrips() { + // An existing header carrying a percent-encoded space must be decoded, then re-encoded back to + // %20 (not corrupted to '+') while a newly merged key is appended. + Map attrs = new LinkedHashMap<>(); + attrs.put("k2", "v2"); + + String merged = OtelBaggageScopeHelper.mergeIntoBaggageHeader("k1=a%20b", attrs); + assertThat(merged).contains("k1=a%20b"); + assertThat(merged).contains("k2=v2"); + assertThat(merged).doesNotContain("+"); + } + + @Test + public void testMergeIntoBaggageHeader_overwriteExistingKey() { + Map attrs = new LinkedHashMap<>(); + attrs.put("k1", "new-v1"); + attrs.put("k3", "v3"); + + String merged = OtelBaggageScopeHelper.mergeIntoBaggageHeader("k1=old-v1,k2=v2", attrs); + // k1 overwritten with new-v1, k2 preserved, k3 appended + assertThat(merged).contains("k1=new-v1"); + assertThat(merged).contains("k2=v2"); + assertThat(merged).contains("k3=v3"); + assertThat(merged).doesNotContain("old-v1"); + } + + @Test + public void testMergeIntoBaggageHeader_existingWithMetadata() { + // W3C baggage allows an optional ; property suffix per entry; parsing must drop the property + Map attrs = new LinkedHashMap<>(); + attrs.put("k2", "v2"); + + String merged = OtelBaggageScopeHelper.mergeIntoBaggageHeader("k1=v1;metadata=x", attrs); + assertThat(merged).contains("k1=v1"); + assertThat(merged).contains("k2=v2"); + assertThat(merged).doesNotContain("metadata"); + } + + @Test + public void testMergeIntoBaggageHeader_nullExisting() { + Map attrs = new LinkedHashMap<>(); + attrs.put("k1", "v1"); + + assertThat(OtelBaggageScopeHelper.mergeIntoBaggageHeader(null, attrs)).isEqualTo("k1=v1"); + assertThat(OtelBaggageScopeHelper.mergeIntoBaggageHeader("", attrs)).isEqualTo("k1=v1"); + } + + @Test + public void testMergeIntoBaggageHeader_emptyAttrs() { + assertThat(OtelBaggageScopeHelper.mergeIntoBaggageHeader("k1=v1", null)).isEqualTo("k1=v1"); + assertThat(OtelBaggageScopeHelper.mergeIntoBaggageHeader("k1=v1", new HashMap<>())).isEqualTo("k1=v1"); + } + + @Test + public void testMergeIntoBaggageHeader_bothEmpty() { + assertThat(OtelBaggageScopeHelper.mergeIntoBaggageHeader(null, null)).isNull(); + assertThat(OtelBaggageScopeHelper.mergeIntoBaggageHeader("", new HashMap<>())).isNull(); + } +}