Skip to content

Commit e3dd43e

Browse files
fix: avoid globally monkeypatching Requests (#1256)
## Summary - remove the import-time `requests_unixsocket.monkeypatch()` call - use a context-managed Unix-socket session for each Edge HSM trust-bundle and signing request - close every session at the end of its operation, independent of client shutdown - preserve existing HTTP behavior without replacing process-wide Requests helpers ## Testing - `uv run --frozen pytest tests/unit/iothub/test_edge_hsm.py tests/unit/iothub/test_sync_clients.py tests/unit/iothub/aio/test_async_inbox.py tests/unit/iothub/aio/test_async_clients.py tests/unit/iothub/test_inbox_manager.py -W error::RuntimeWarning -q` (1564 passed, 6 skipped) - regression test reloads the Edge HSM module and verifies `requests_unixsocket.monkeypatch()` is not called - tests verify both HSM operations exit their session context - targeted Ruff and Black checks --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent 8c0d51a commit e3dd43e

2 files changed

Lines changed: 117 additions & 77 deletions

File tree

‎azure-iot-device/azure/iot/device/iothub/edge_hsm.py‎

Lines changed: 42 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
from azure.iot.device.common.auth.signing_mechanism import SigningMechanism
1414
from azure.iot.device import user_agent
1515

16-
requests_unixsocket.monkeypatch()
1716
logger = logging.getLogger(__name__)
1817

1918

@@ -58,27 +57,27 @@ def get_certificate(self):
5857
5958
:raises: IoTEdgeError if unable to retrieve the certificate.
6059
"""
61-
r = requests.get(
62-
self.workload_uri + "trust-bundle",
63-
params={"api-version": self.api_version},
64-
headers={"User-Agent": urllib.parse.quote_plus(user_agent.get_iothub_user_agent())},
65-
)
66-
# Validate that the request was successful
67-
try:
68-
r.raise_for_status()
69-
except requests.exceptions.HTTPError as e:
70-
raise IoTEdgeError("Unable to get trust bundle from Edge") from e
71-
# Decode the trust bundle
72-
try:
73-
bundle = r.json()
74-
except ValueError as e:
75-
raise IoTEdgeError("Unable to decode trust bundle") from e
76-
# Retrieve the certificate
77-
try:
78-
cert = bundle["certificate"]
79-
except KeyError as e:
80-
raise IoTEdgeError("No certificate in trust bundle") from e
81-
return cert
60+
with requests_unixsocket.Session() as session:
61+
r = session.get(
62+
self.workload_uri + "trust-bundle",
63+
params={"api-version": self.api_version},
64+
headers={"User-Agent": urllib.parse.quote_plus(user_agent.get_iothub_user_agent())},
65+
)
66+
# Validate that the request was successful
67+
try:
68+
r.raise_for_status()
69+
except requests.exceptions.HTTPError as e:
70+
raise IoTEdgeError("Unable to get trust bundle from Edge") from e
71+
# Decode the trust bundle
72+
try:
73+
bundle = r.json()
74+
except ValueError as e:
75+
raise IoTEdgeError("Unable to decode trust bundle") from e
76+
# Retrieve the certificate
77+
try:
78+
return bundle["certificate"]
79+
except KeyError as e:
80+
raise IoTEdgeError("No certificate in trust bundle") from e
8281

8382
def sign(self, data_str):
8483
"""
@@ -99,26 +98,27 @@ def sign(self, data_str):
9998
)
10099
sign_request = {"keyId": "primary", "algo": "HMACSHA256", "data": encoded_data_str}
101100

102-
r = requests.post( # can we use json field instead of data?
103-
url=path,
104-
params={"api-version": self.api_version},
105-
headers={"User-Agent": urllib.parse.quote(user_agent.get_iothub_user_agent(), safe="")},
106-
data=json.dumps(sign_request),
107-
)
108-
try:
109-
r.raise_for_status()
110-
except requests.exceptions.HTTPError as e:
111-
raise IoTEdgeError("Unable to sign data") from e
112-
try:
113-
sign_response = r.json()
114-
except ValueError as e:
115-
raise IoTEdgeError("Unable to decode signed data") from e
116-
try:
117-
signed_data_str = sign_response["digest"]
118-
except KeyError as e:
119-
raise IoTEdgeError("No signed data received") from e
120-
121-
return signed_data_str # what format is this? string? bytes?
101+
with requests_unixsocket.Session() as session:
102+
r = session.post( # can we use json field instead of data?
103+
url=path,
104+
params={"api-version": self.api_version},
105+
headers={
106+
"User-Agent": urllib.parse.quote(user_agent.get_iothub_user_agent(), safe="")
107+
},
108+
data=json.dumps(sign_request),
109+
)
110+
try:
111+
r.raise_for_status()
112+
except requests.exceptions.HTTPError as e:
113+
raise IoTEdgeError("Unable to sign data") from e
114+
try:
115+
sign_response = r.json()
116+
except ValueError as e:
117+
raise IoTEdgeError("Unable to decode signed data") from e
118+
try:
119+
return sign_response["digest"]
120+
except KeyError as e:
121+
raise IoTEdgeError("No signed data received") from e
122122

123123

124124
def _format_socket_uri(old_uri):

‎tests/unit/iothub/test_edge_hsm.py‎

Lines changed: 75 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -4,39 +4,59 @@
44
# license information.
55
# --------------------------------------------------------------------------
66

7+
import importlib
78
import pytest
89
import logging
910
import requests
11+
import requests_unixsocket
1012
import json
1113
import base64
1214
import urllib
13-
from azure.iot.device.iothub.edge_hsm import IoTEdgeHsm, IoTEdgeError
14-
from azure.iot.device import user_agent
1515

16+
# Keep module-qualified references because the import regression test reloads this module.
17+
# Directly imported classes would retain their pre-reload identities.
18+
from azure.iot.device.iothub import edge_hsm as edge_hsm_module
19+
from azure.iot.device import user_agent
1620

1721
logging.basicConfig(level=logging.DEBUG)
1822

1923

2024
@pytest.fixture
2125
def edge_hsm():
22-
return IoTEdgeHsm(
26+
return edge_hsm_module.IoTEdgeHsm(
2327
module_id="my_module_id",
2428
generation_id="module_generation_id",
2529
workload_uri="unix:///var/run/iotedge/workload.sock",
2630
api_version="my_api_version",
2731
)
2832

2933

34+
@pytest.fixture
35+
def mock_unix_session(mocker):
36+
mock_session_constructor = mocker.patch.object(requests_unixsocket, "Session")
37+
mock_session = mock_session_constructor.return_value
38+
mock_session.__enter__.return_value = mock_session
39+
return mock_session
40+
41+
3042
@pytest.mark.describe("IoTEdgeHsm - Instantiation")
3143
class TestIoTEdgeHsmInstantiation(object):
44+
@pytest.mark.it("Does not monkeypatch the global requests API when imported")
45+
def test_does_not_monkeypatch_requests(self, mocker):
46+
mock_monkeypatch = mocker.patch.object(requests_unixsocket, "monkeypatch")
47+
48+
importlib.reload(edge_hsm_module)
49+
50+
assert mock_monkeypatch.call_count == 0
51+
3252
@pytest.mark.it("URL encodes the provided module_id parameter and sets it as an attribute")
3353
def test_encode_and_set_module_id(self):
3454
module_id = "my_module_id"
3555
generation_id = "my_generation_id"
3656
api_version = "my_api_version"
3757
workload_uri = "unix:///var/run/iotedge/workload.sock"
3858

39-
edge_hsm = IoTEdgeHsm(
59+
edge_hsm = edge_hsm_module.IoTEdgeHsm(
4060
module_id=module_id,
4161
generation_id=generation_id,
4262
workload_uri=workload_uri,
@@ -64,7 +84,7 @@ def test_workload_uri_formatting(self, workload_uri, expected_formatted_uri):
6484
generation_id = "my_generation_id"
6585
api_version = "my_api_version"
6686

67-
edge_hsm = IoTEdgeHsm(
87+
edge_hsm = edge_hsm_module.IoTEdgeHsm(
6888
module_id=module_id,
6989
generation_id=generation_id,
7090
workload_uri=workload_uri,
@@ -80,7 +100,7 @@ def test_set_generation_id(self):
80100
api_version = "my_api_version"
81101
workload_uri = "unix:///var/run/iotedge/workload.sock"
82102

83-
edge_hsm = IoTEdgeHsm(
103+
edge_hsm = edge_hsm_module.IoTEdgeHsm(
84104
module_id=module_id,
85105
generation_id=generation_id,
86106
workload_uri=workload_uri,
@@ -96,7 +116,7 @@ def test_set_api_version(self):
96116
api_version = "my_api_version"
97117
workload_uri = "unix:///var/run/iotedge/workload.sock"
98118

99-
edge_hsm = IoTEdgeHsm(
119+
edge_hsm = edge_hsm_module.IoTEdgeHsm(
100120
module_id=module_id,
101121
generation_id=generation_id,
102122
workload_uri=workload_uri,
@@ -108,9 +128,19 @@ def test_set_api_version(self):
108128

109129
@pytest.mark.describe("IoTEdgeHsm - .get_certificate()")
110130
class TestIoTEdgeHsmGetCertificate(object):
131+
@pytest.mark.it("Closes the Unix socket session after the request")
132+
def test_closes_session(self, mocker, edge_hsm, mock_unix_session):
133+
mock_unix_session.get.return_value.json.return_value = {"certificate": "my certificate"}
134+
135+
edge_hsm.get_certificate()
136+
137+
assert requests_unixsocket.Session.call_args == mocker.call()
138+
assert mock_unix_session.__enter__.call_args == mocker.call()
139+
assert mock_unix_session.__exit__.call_count == 1
140+
111141
@pytest.mark.it("Sends an HTTP GET request to retrieve the trust bundle from Edge")
112-
def test_requests_trust_bundle(self, mocker, edge_hsm):
113-
mock_request_get = mocker.patch.object(requests, "get")
142+
def test_requests_trust_bundle(self, mocker, edge_hsm, mock_unix_session):
143+
mock_request_get = mock_unix_session.get
114144
expected_url = edge_hsm.workload_uri + "trust-bundle"
115145
expected_params = {"api-version": edge_hsm.api_version}
116146
expected_headers = {
@@ -125,8 +155,8 @@ def test_requests_trust_bundle(self, mocker, edge_hsm):
125155
)
126156

127157
@pytest.mark.it("Returns the certificate from the trust bundle received from Edge")
128-
def test_returns_certificate(self, mocker, edge_hsm):
129-
mock_request_get = mocker.patch.object(requests, "get")
158+
def test_returns_certificate(self, edge_hsm, mock_unix_session):
159+
mock_request_get = mock_unix_session.get
130160
mock_response = mock_request_get.return_value
131161
certificate = "my certificate"
132162
mock_response.json.return_value = {"certificate": certificate}
@@ -136,47 +166,57 @@ def test_returns_certificate(self, mocker, edge_hsm):
136166
assert returned_cert is certificate
137167

138168
@pytest.mark.it("Raises IoTEdgeError if a bad request is made to Edge")
139-
def test_bad_request(self, mocker, edge_hsm):
140-
mock_request_get = mocker.patch.object(requests, "get")
169+
def test_bad_request(self, edge_hsm, mock_unix_session):
170+
mock_request_get = mock_unix_session.get
141171
mock_response = mock_request_get.return_value
142172
error = requests.exceptions.HTTPError()
143173
mock_response.raise_for_status.side_effect = error
144174

145-
with pytest.raises(IoTEdgeError) as e_info:
175+
with pytest.raises(edge_hsm_module.IoTEdgeError) as e_info:
146176
edge_hsm.get_certificate()
147177
assert e_info.value.__cause__ is error
148178

149179
@pytest.mark.it("Raises IoTEdgeError if there is an error in json decoding the trust bundle")
150-
def test_bad_json(self, mocker, edge_hsm):
151-
mock_request_get = mocker.patch.object(requests, "get")
180+
def test_bad_json(self, edge_hsm, mock_unix_session):
181+
mock_request_get = mock_unix_session.get
152182
mock_response = mock_request_get.return_value
153183
error = ValueError()
154184
mock_response.json.side_effect = error
155185

156-
with pytest.raises(IoTEdgeError) as e_info:
186+
with pytest.raises(edge_hsm_module.IoTEdgeError) as e_info:
157187
edge_hsm.get_certificate()
158188
assert e_info.value.__cause__ is error
159189

160190
@pytest.mark.it("Raises IoTEdgeError if the certificate is missing from the trust bundle")
161-
def test_bad_trust_bundle(self, mocker, edge_hsm):
162-
mock_request_get = mocker.patch.object(requests, "get")
191+
def test_bad_trust_bundle(self, edge_hsm, mock_unix_session):
192+
mock_request_get = mock_unix_session.get
163193
mock_response = mock_request_get.return_value
164194
# Return an empty json dict with no 'certificate' key
165195
mock_response.json.return_value = {}
166196

167-
with pytest.raises(IoTEdgeError):
197+
with pytest.raises(edge_hsm_module.IoTEdgeError):
168198
edge_hsm.get_certificate()
169199

170200

171201
@pytest.mark.describe("IoTEdgeHsm - .sign()")
172202
class TestIoTEdgeHsmSign(object):
203+
@pytest.mark.it("Closes the Unix socket session after the request")
204+
def test_closes_session(self, mocker, edge_hsm, mock_unix_session):
205+
mock_unix_session.post.return_value.json.return_value = {"digest": "somedigest"}
206+
207+
edge_hsm.sign("somedata")
208+
209+
assert requests_unixsocket.Session.call_args == mocker.call()
210+
assert mock_unix_session.__enter__.call_args == mocker.call()
211+
assert mock_unix_session.__exit__.call_count == 1
212+
173213
@pytest.mark.it(
174214
"Makes an HTTP request to Edge to sign a piece of string data using the HMAC-SHA256 algorithm"
175215
)
176-
def test_requests_data_signing(self, mocker, edge_hsm):
216+
def test_requests_data_signing(self, mocker, edge_hsm, mock_unix_session):
177217
data_str = "somedata"
178218
data_str_b64 = "c29tZWRhdGE="
179-
mock_request_post = mocker.patch.object(requests, "post")
219+
mock_request_post = mock_unix_session.post
180220
mock_request_post.return_value.json.return_value = {"digest": "somedigest"}
181221
expected_url = "{workload_uri}modules/{module_id}/genid/{generation_id}/sign".format(
182222
workload_uri=edge_hsm.workload_uri,
@@ -197,12 +237,12 @@ def test_requests_data_signing(self, mocker, edge_hsm):
197237
)
198238

199239
@pytest.mark.it("Base64 encodes the string data in the request")
200-
def test_b64_encodes_data(self, mocker, edge_hsm):
240+
def test_b64_encodes_data(self, edge_hsm, mock_unix_session):
201241
# This test is actually implicitly tested in the first test, but it's
202242
# important to have an explicit test for it since it's a requirement
203243
data_str = "somedata"
204244
data_str_b64 = base64.b64encode(data_str.encode("utf-8")).decode()
205-
mock_request_post = mocker.patch.object(requests, "post")
245+
mock_request_post = mock_unix_session.post
206246
mock_request_post.return_value.json.return_value = {"digest": "somedigest"}
207247

208248
edge_hsm.sign(data_str)
@@ -213,41 +253,41 @@ def test_b64_encodes_data(self, mocker, edge_hsm):
213253
assert sent_data == data_str_b64
214254

215255
@pytest.mark.it("Returns the signed data received from Edge")
216-
def test_returns_signed_data(self, mocker, edge_hsm):
256+
def test_returns_signed_data(self, edge_hsm, mock_unix_session):
217257
expected_digest = "somedigest"
218-
mock_request_post = mocker.patch.object(requests, "post")
258+
mock_request_post = mock_unix_session.post
219259
mock_request_post.return_value.json.return_value = {"digest": expected_digest}
220260

221261
signed_data = edge_hsm.sign("somedata")
222262

223263
assert signed_data == expected_digest
224264

225265
@pytest.mark.it("Raises IoTEdgeError if a bad request is made to EdgeHub")
226-
def test_bad_request(self, mocker, edge_hsm):
227-
mock_request_post = mocker.patch.object(requests, "post")
266+
def test_bad_request(self, edge_hsm, mock_unix_session):
267+
mock_request_post = mock_unix_session.post
228268
mock_response = mock_request_post.return_value
229269
error = requests.exceptions.HTTPError()
230270
mock_response.raise_for_status.side_effect = error
231271

232-
with pytest.raises(IoTEdgeError) as e_info:
272+
with pytest.raises(edge_hsm_module.IoTEdgeError) as e_info:
233273
edge_hsm.sign("somedata")
234274
assert e_info.value.__cause__ is error
235275

236276
@pytest.mark.it("Raises IoTEdgeError if there is an error in json decoding the signed response")
237-
def test_bad_json(self, mocker, edge_hsm):
238-
mock_request_post = mocker.patch.object(requests, "post")
277+
def test_bad_json(self, edge_hsm, mock_unix_session):
278+
mock_request_post = mock_unix_session.post
239279
mock_response = mock_request_post.return_value
240280
error = ValueError()
241281
mock_response.json.side_effect = error
242-
with pytest.raises(IoTEdgeError) as e_info:
282+
with pytest.raises(edge_hsm_module.IoTEdgeError) as e_info:
243283
edge_hsm.sign("somedata")
244284
assert e_info.value.__cause__ is error
245285

246286
@pytest.mark.it("Raises IoTEdgeError if the signed data is missing from the response")
247-
def test_bad_response(self, mocker, edge_hsm):
248-
mock_request_post = mocker.patch.object(requests, "post")
287+
def test_bad_response(self, edge_hsm, mock_unix_session):
288+
mock_request_post = mock_unix_session.post
249289
mock_response = mock_request_post.return_value
250290
mock_response.json.return_value = {}
251291

252-
with pytest.raises(IoTEdgeError):
292+
with pytest.raises(edge_hsm_module.IoTEdgeError):
253293
edge_hsm.sign("somedata")

0 commit comments

Comments
 (0)