@@ -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
356389And now, to plug into the client code:
0 commit comments