From d0bde24a338bd84ec685f1f0f47117779f709cfc Mon Sep 17 00:00:00 2001
From: Fools <54661071+Charlie-lizhihan@users.noreply.github.com>
Date: Thu, 13 Jun 2024 16:56:09 +0100
Subject: [PATCH] bugfix: RandomIdGenerator can generate invalid Span/Trace Ids
 (#3949)

---
 CHANGELOG.md                                  |  2 ++
 .../opentelemetry/sdk/trace/id_generator.py   | 12 +++++--
 opentelemetry-sdk/tests/trace/test_trace.py   | 32 +++++++++++++++++++
 3 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2c495129ea7..68b9a47006b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+- Fix RandomIdGenerator can generate invalid Span/Trace Ids
+  ([#3949](https://github.com/open-telemetry/opentelemetry-python/pull/3949))
 - Add Python 3.12 to tox
   ([#3616](https://github.com/open-telemetry/opentelemetry-python/pull/3616))
 
diff --git a/opentelemetry-sdk/src/opentelemetry/sdk/trace/id_generator.py b/opentelemetry-sdk/src/opentelemetry/sdk/trace/id_generator.py
index 62b12a94921..cd1f89bcde2 100644
--- a/opentelemetry-sdk/src/opentelemetry/sdk/trace/id_generator.py
+++ b/opentelemetry-sdk/src/opentelemetry/sdk/trace/id_generator.py
@@ -15,6 +15,8 @@
 import abc
 import random
 
+from opentelemetry import trace
+
 
 class IdGenerator(abc.ABC):
     @abc.abstractmethod
@@ -46,7 +48,13 @@ class RandomIdGenerator(IdGenerator):
     """
 
     def generate_span_id(self) -> int:
-        return random.getrandbits(64)
+        span_id = random.getrandbits(64)
+        while span_id == trace.INVALID_SPAN_ID:
+            span_id = random.getrandbits(64)
+        return span_id
 
     def generate_trace_id(self) -> int:
-        return random.getrandbits(128)
+        trace_id = random.getrandbits(128)
+        while trace_id == trace.INVALID_TRACE_ID:
+            trace_id = random.getrandbits(128)
+        return trace_id
diff --git a/opentelemetry-sdk/tests/trace/test_trace.py b/opentelemetry-sdk/tests/trace/test_trace.py
index 30f4f0e2731..2a064be80f1 100644
--- a/opentelemetry-sdk/tests/trace/test_trace.py
+++ b/opentelemetry-sdk/tests/trace/test_trace.py
@@ -2061,3 +2061,35 @@ def test_tracer_provider_init_default(self, resource_patch, sample_patch):
         sample_patch.assert_called_once()
         self.assertIsNotNone(tracer_provider._span_limits)
         self.assertIsNotNone(tracer_provider._atexit_handler)
+
+
+class TestRandomIdGenerator(unittest.TestCase):
+    _TRACE_ID_MAX_VALUE = 2**128 - 1
+    _SPAN_ID_MAX_VALUE = 2**64 - 1
+
+    @patch(
+        "random.getrandbits",
+        side_effect=[trace_api.INVALID_SPAN_ID, 0x00000000DEADBEF0],
+    )
+    def test_generate_span_id_avoids_invalid(self, mock_getrandbits):
+        generator = RandomIdGenerator()
+        span_id = generator.generate_span_id()
+
+        self.assertNotEqual(span_id, trace_api.INVALID_SPAN_ID)
+        mock_getrandbits.assert_any_call(64)
+        self.assertEqual(mock_getrandbits.call_count, 2)
+
+    @patch(
+        "random.getrandbits",
+        side_effect=[
+            trace_api.INVALID_TRACE_ID,
+            0x000000000000000000000000DEADBEEF,
+        ],
+    )
+    def test_generate_trace_id_avoids_invalid(self, mock_getrandbits):
+        generator = RandomIdGenerator()
+        trace_id = generator.generate_trace_id()
+
+        self.assertNotEqual(trace_id, trace_api.INVALID_TRACE_ID)
+        mock_getrandbits.assert_any_call(128)
+        self.assertEqual(mock_getrandbits.call_count, 2)