From bcc97a2f8b84072210094d7877bf421a5224e028 Mon Sep 17 00:00:00 2001
From: Adrian Cole <adriancole@users.noreply.github.com>
Date: Wed, 21 Mar 2018 21:57:34 +0800
Subject: [PATCH] Fixes bug where debug did not imply sampled (#663)

It is true the b3-propagation spec suggests debug implies sampled.
This leniently parses a missing `X-B3-Sampled` header as long as
`X-B3-Flags: 1` is present.

Fixes #661
---
 .../java/brave/propagation/B3PropagationTest.java  | 12 ++++++++++++
 .../java/brave/propagation/TraceIdContext.java     |  3 +++
 .../java/brave/propagation/TraceIdContextTest.java | 14 ++++++++++++++
 3 files changed, 29 insertions(+)

diff --git a/brave-tests/src/test/java/brave/propagation/B3PropagationTest.java b/brave-tests/src/test/java/brave/propagation/B3PropagationTest.java
index 4c7432f797..33968ecc66 100644
--- a/brave-tests/src/test/java/brave/propagation/B3PropagationTest.java
+++ b/brave-tests/src/test/java/brave/propagation/B3PropagationTest.java
@@ -72,4 +72,16 @@ public class B3PropagationTest extends PropagationTest<String> {
     assertThat(result)
         .isEqualTo(SamplingFlags.EMPTY);
   }
+
+  @Test public void extractTraceContext_debug_with_ids() {
+    MapEntry mapEntry = new MapEntry();
+    map.put("X-B3-TraceId", "463ac35c9f6413ad48485a3953bb6124"); // ok
+    map.put("X-B3-SpanId", "48485a3953bb6124"); // ok
+    map.put("X-B3-Flags", "1"); // accidentally missing sampled flag
+
+    TraceContext result = propagation().extractor(mapEntry).extract(map).context();
+
+    assertThat(result.sampled())
+        .isTrue();
+  }
 }
diff --git a/brave/src/main/java/brave/propagation/TraceIdContext.java b/brave/src/main/java/brave/propagation/TraceIdContext.java
index 06ec72fa64..499b54d2f3 100644
--- a/brave/src/main/java/brave/propagation/TraceIdContext.java
+++ b/brave/src/main/java/brave/propagation/TraceIdContext.java
@@ -224,9 +224,12 @@ InternalBuilder sampled(@Nullable Boolean sampled) {
       return this;
     }
 
+    /** Ensures sampled is set when debug is */
     InternalBuilder debug(boolean debug) {
       if (debug) {
         flags |= FLAG_DEBUG;
+        flags |= FLAG_SAMPLED_SET;
+        flags |= FLAG_SAMPLED;
       } else {
         flags &= ~FLAG_DEBUG;
       }
diff --git a/brave/src/test/java/brave/propagation/TraceIdContextTest.java b/brave/src/test/java/brave/propagation/TraceIdContextTest.java
index c30e2af458..4535c829aa 100644
--- a/brave/src/test/java/brave/propagation/TraceIdContextTest.java
+++ b/brave/src/test/java/brave/propagation/TraceIdContextTest.java
@@ -32,6 +32,20 @@ public class TraceIdContextTest {
         .isEqualTo("00000000000000de000000000000014d");
   }
 
+  @Test public void debugImpliesSampled() {
+    TraceIdContext primitives = context.toBuilder()
+        .debug(true)
+        .build();
+
+    TraceIdContext objects = context.toBuilder()
+        .sampled(Boolean.TRUE)
+        .debug(Boolean.TRUE)
+        .build();
+
+    assertThat(primitives)
+        .isEqualToComparingFieldByField(objects);
+  }
+
   @Test public void canUsePrimitiveOverloads() {
     TraceIdContext primitives = context.toBuilder()
         .sampled(true)