Skip to content

Commit eb474e5

Browse files
jkebingerclaude
andauthored
Prevent processing of zero-byte config payloads (#12)
Added validation to reject zero-byte config payloads from both HTTP and SSE endpoints, treating them as connection errors that trigger retries. Changes: - Added zero-byte validation in config_sdk.py for HTTP responses - Added zero-byte validation in _sse_connection_manager.py for SSE events - Return False/early when zero-byte payloads detected to trigger retry logic - Added comprehensive tests for both HTTP and SSE scenarios This prevents the SDK from attempting to process empty config data which would likely cause parsing errors or unexpected behavior. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-authored-by: Claude <noreply@anthropic.com>
1 parent b02d106 commit eb474e5

File tree

3 files changed

+185
-1
lines changed

3 files changed

+185
-1
lines changed

sdk_reforge/_sse_connection_manager.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,14 @@ def process_response(self, response: Response) -> None:
9797
logger.info("Client is shutting down, exiting SSE event loop")
9898
return
9999
if event.data:
100-
configs = Prefab.Configs.FromString(base64.b64decode(event.data))
100+
decoded_data = base64.b64decode(event.data)
101+
if not decoded_data or len(decoded_data) == 0:
102+
logger.warning(
103+
"Received zero-byte config payload from SSE stream, treating as connection error"
104+
)
105+
# Return early to trigger reconnection logic
106+
return
107+
configs = Prefab.Configs.FromString(decoded_data)
101108
self.config_client.load_configs(configs, "sse_streaming")
102109
self.sse_client.close()
103110
self.sse_client = None

sdk_reforge/config_sdk.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,6 +165,11 @@ def load_checkpoint_from_api_cdn(self):
165165
allow_cache=True,
166166
)
167167
if response.ok:
168+
if not response.content or len(response.content) == 0:
169+
logger.warning(
170+
"Received zero-byte config payload from remote_cdn_api, treating as connection error"
171+
)
172+
return False
168173
configs = Prefab.Configs.FromString(response.content)
169174
self.load_configs(configs, "remote_api_cdn")
170175
return True

tests/test_zero_byte_configs.py

Lines changed: 172 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,172 @@
1+
from __future__ import annotations
2+
3+
import unittest
4+
from unittest.mock import Mock, patch, MagicMock
5+
import base64
6+
7+
from sdk_reforge.config_sdk import ConfigSDK
8+
from sdk_reforge._sse_connection_manager import SSEConnectionManager
9+
import prefab_pb2 as Prefab
10+
11+
12+
class TestZeroByteConfigHandling(unittest.TestCase):
13+
@patch("sdk_reforge.config_sdk.logger")
14+
def test_http_config_zero_byte_payload(self, mock_logger: MagicMock) -> None:
15+
"""Test that zero-byte HTTP config responses are treated as errors"""
16+
# Create mock objects
17+
mock_api_client = Mock()
18+
mock_response = Mock()
19+
mock_response.ok = True
20+
mock_response.content = b"" # Zero-byte response
21+
mock_api_client.resilient_request.return_value = mock_response
22+
23+
# Create ConfigSDK with mocked dependencies
24+
options = Mock()
25+
options.api_key = "test_key"
26+
config_sdk = ConfigSDK(options)
27+
config_sdk.api_client = mock_api_client
28+
config_sdk.config_loader = Mock()
29+
config_sdk.config_loader.highwater_mark = 123
30+
31+
# Test that load_checkpoint_from_api_cdn returns False for zero-byte response
32+
result = config_sdk.load_checkpoint_from_api_cdn()
33+
34+
self.assertFalse(result)
35+
mock_logger.warning.assert_called_with(
36+
"Received zero-byte config payload from remote_cdn_api, treating as connection error"
37+
)
38+
39+
@patch("sdk_reforge.config_sdk.logger")
40+
def test_http_config_valid_payload(self, mock_logger: MagicMock) -> None:
41+
"""Test that valid HTTP config responses are processed normally"""
42+
# Create mock objects
43+
mock_api_client = Mock()
44+
mock_response = Mock()
45+
mock_response.ok = True
46+
47+
# Create a valid Prefab.Configs message
48+
configs = Prefab.Configs()
49+
config = configs.configs.add()
50+
config.key = "test_key"
51+
valid_content = configs.SerializeToString()
52+
mock_response.content = valid_content
53+
54+
mock_api_client.resilient_request.return_value = mock_response
55+
56+
# Create ConfigSDK with mocked dependencies
57+
options = Mock()
58+
options.api_key = "test_key"
59+
config_sdk = ConfigSDK(options)
60+
config_sdk.api_client = mock_api_client
61+
config_sdk.config_loader = Mock()
62+
config_sdk.config_loader.highwater_mark = 123
63+
config_sdk.load_configs = Mock()
64+
65+
# Test that load_checkpoint_from_api_cdn returns True for valid response
66+
result = config_sdk.load_checkpoint_from_api_cdn()
67+
68+
self.assertTrue(result)
69+
config_sdk.load_configs.assert_called_once()
70+
# Check that the specific zero-byte warning was not called
71+
warning_calls = [str(call) for call in mock_logger.warning.call_args_list]
72+
self.assertNotIn(
73+
"call('Received zero-byte config payload from remote_cdn_api, treating as connection error')",
74+
warning_calls,
75+
)
76+
77+
def test_sse_config_zero_byte_payload(self) -> None:
78+
"""Test that zero-byte SSE config payloads trigger reconnection"""
79+
# Create mock objects
80+
mock_api_client = Mock()
81+
mock_config_client = Mock()
82+
mock_config_client.is_shutting_down.return_value = False
83+
84+
# Create SSEConnectionManager
85+
sse_manager = SSEConnectionManager(
86+
mock_api_client, mock_config_client, ["https://test.example.com"]
87+
)
88+
89+
# Mock SSE event with a base64 string that will decode to empty bytes
90+
# We need a non-empty string that produces empty bytes when decoded
91+
# In practice, this would happen when the server sends an empty base64 string
92+
# For testing, we'll use a mock that bypasses base64 entirely
93+
mock_event = Mock()
94+
mock_event.data = "dummy" # Non-empty string to pass the if check
95+
96+
mock_sse_client = Mock()
97+
mock_sse_client.events.return_value = iter([mock_event])
98+
mock_sse_client.close = Mock()
99+
100+
mock_response = Mock()
101+
102+
with patch(
103+
"sdk_reforge._sse_connection_manager.sseclient.SSEClient",
104+
return_value=mock_sse_client,
105+
):
106+
with patch(
107+
"sdk_reforge._sse_connection_manager.base64.b64decode"
108+
) as mock_b64decode:
109+
mock_b64decode.return_value = b"" # Return empty bytes
110+
with patch("sdk_reforge._sse_connection_manager.logger") as mock_logger:
111+
# The process_response method should return early on zero-byte payload
112+
sse_manager.process_response(mock_response)
113+
114+
# Verify that warning was logged and configs were not loaded
115+
mock_logger.warning.assert_called_with(
116+
"Received zero-byte config payload from SSE stream, treating as connection error"
117+
)
118+
119+
mock_config_client.load_configs.assert_not_called()
120+
# SSE client should not be closed since we returned early
121+
mock_sse_client.close.assert_not_called()
122+
123+
@patch("sdk_reforge._sse_connection_manager.logger")
124+
def test_sse_config_valid_payload(self, mock_logger: MagicMock) -> None:
125+
"""Test that valid SSE config payloads are processed normally"""
126+
# Create mock objects
127+
mock_api_client = Mock()
128+
mock_config_client = Mock()
129+
mock_config_client.is_shutting_down.return_value = False
130+
131+
# Create SSEConnectionManager
132+
sse_manager = SSEConnectionManager(
133+
mock_api_client, mock_config_client, ["https://test.example.com"]
134+
)
135+
136+
# Create a valid Prefab.Configs message
137+
configs = Prefab.Configs()
138+
config = configs.configs.add()
139+
config.key = "test_key"
140+
valid_content = configs.SerializeToString()
141+
142+
# Mock SSE event with valid payload
143+
mock_event = Mock()
144+
mock_event.data = base64.b64encode(valid_content).decode("utf-8")
145+
146+
mock_sse_client = Mock()
147+
mock_sse_client.events.return_value = iter([mock_event])
148+
mock_sse_client.close = Mock()
149+
150+
mock_response = Mock()
151+
152+
with patch(
153+
"sdk_reforge._sse_connection_manager.sseclient.SSEClient",
154+
return_value=mock_sse_client,
155+
):
156+
with patch(
157+
"sdk_reforge._sse_connection_manager.Prefab.Configs.FromString"
158+
) as mock_from_string:
159+
mock_from_string.return_value = configs
160+
sse_manager.process_response(mock_response)
161+
162+
# Verify that configs were loaded and no warning was logged
163+
mock_config_client.load_configs.assert_called_once_with(
164+
configs, "sse_streaming"
165+
)
166+
mock_logger.warning.assert_not_called()
167+
# SSE client should be closed after successful processing
168+
mock_sse_client.close.assert_called_once()
169+
170+
171+
if __name__ == "__main__":
172+
unittest.main()

0 commit comments

Comments
 (0)