[RPC][Tracker] Reject invalid tracker message sizes and consume frame header#19591
[RPC][Tracker] Reject invalid tracker message sizes and consume frame header#19591cchung100m wants to merge 2 commits into
Conversation
There was a problem hiding this comment.
Code Review
This pull request enhances the RPC tracker's robustness by introducing a maximum message size limit (MAX_TRACKER_MSG_BYTES) and refactoring the message processing loop in python/tvm/rpc/tracker.py to validate incoming data. The review feedback suggests improving the loop termination condition by checking the socket state and recommends wrapping the JSON parsing and handler execution in a try...except block to prevent crashes from malformed inputs.
| @@ -222,18 +224,21 @@ def on_message(self, message): | |||
|
|
|||
| while True: | |||
There was a problem hiding this comment.
The while True loop continues processing messages from the internal buffer even if the connection is closed during an iteration (e.g., by self.close() in the previous iteration or inside call_handler). Changing this to while self._sock: ensures that the loop terminates immediately if the connection is no longer active.
| while True: | |
| while self._sock: |
|
Close PR because the original issue is tracked by #19586 |
Hi Committers,
This PR fixes #19585.
Root Cause
TCPEventHandler.on_messageparsed the 4-byte int 32 length header directly from the accumulated buffer without limits and left the header bytes in the buffer until the full payload arrived.Solution
MAX_TRACKER_MSG_BYTES = 1 << 20(1MiB).