Skip to content

Commit a65d727

Browse files
authored
update LRO customization guidance (Azure#44782)
1 parent 53b9269 commit a65d727

1 file changed

Lines changed: 41 additions & 8 deletions

File tree

doc/dev/customize_long_running_operation.md

Lines changed: 41 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -323,10 +323,32 @@ class CustomPollingMethod(PollingMethod):
323323
324324
In standard LROs, the PipelineResponse is serialized here, however, there may be a need to
325325
customize this further depending on your scenario.
326-
"""
327-
import pickle
328326
329-
return base64.b64encode(pickle.dumps(self._initial_response)).decode("ascii")
327+
It's recommended to:
328+
1. Include a version number for future compatibility.
329+
2. Filter headers to only include those needed for LRO rehydration, avoiding sensitive data.
330+
"""
331+
import json
332+
333+
# Headers needed for LRO rehydration - use an allowlist approach for security
334+
lro_headers = {"operation-location", "location", "content-type", "retry-after"}
335+
response = self._initial_response.http_response
336+
filtered_headers = {k: v for k, v in response.headers.items() if k.lower() in lro_headers}
337+
338+
# Use a versioned token schema: {"version": <int>, "data": <your state>}
339+
# This allows for future compatibility checking when deserializing
340+
token = {
341+
"version": 1,
342+
"data": {
343+
"file_id": self.file_id,
344+
"response": {
345+
"status_code": response.status_code,
346+
"headers": filtered_headers,
347+
"content": base64.b64encode(response.content).decode("ascii"),
348+
},
349+
},
350+
}
351+
return base64.b64encode(json.dumps(token).encode("utf-8")).decode("ascii")
330352

331353
@classmethod
332354
def from_continuation_token(cls, continuation_token: str, **kwargs: Any) -> Tuple[Any, PipelineResponse, Callable]:
@@ -345,12 +367,23 @@ class CustomPollingMethod(PollingMethod):
345367
"Need kwarg 'deserialization_callback' to be recreated from continuation_token"
346368
)
347369

348-
import pickle
370+
import json
371+
372+
token = json.loads(base64.b64decode(continuation_token).decode("utf-8"))
373+
374+
# Validate token schema and version for compatibility
375+
if not isinstance(token, dict) or "version" not in token or "data" not in token:
376+
raise ValueError("Invalid continuation token format.")
377+
if token["version"] != 1:
378+
raise ValueError(
379+
"This continuation token is not compatible with this version. "
380+
"It may have been generated by a different version."
381+
)
349382

350-
initial_response = pickle.loads(base64.b64decode(continuation_token)) # nosec
351-
# Restore the transport in the context
352-
initial_response.context.transport = client._client._pipeline._transport # pylint: disable=protected-access
353-
return client, initial_response, deserialization_callback
383+
# Extract the state from the "data" field
384+
state = token["data"]
385+
# The file_id and other state can be extracted and used to resume polling
386+
return client, state, deserialization_callback
354387
```
355388

356389
And now, to plug into the client code:

0 commit comments

Comments
 (0)