gemini-code-assist[bot] commented on code in PR #19591:
URL: https://github.com/apache/tvm/pull/19591#discussion_r3274683055
##########
python/tvm/rpc/tracker.py:
##########
@@ -222,18 +224,21 @@ def on_message(self, message):
while True:
if self._msg_size == 0:
- if len(self._data) >= 4:
- self._msg_size = struct.unpack("<i", self._data[:4])[0]
- else:
+ if len(self._data) < 4:
return
- if self._msg_size != 0 and len(self._data) >= self._msg_size + 4:
- msg = py_str(bytes(self._data[4 : 4 + self._msg_size]))
- del self._data[: 4 + self._msg_size]
- self._msg_size = 0
- # pylint: disable=broad-except
- self.call_handler(json.loads(msg))
- else:
+ self._msg_size = struct.unpack("<i", self._data[:4])[0]
+ del self._data[:4]
+ if self._msg_size <= 0 or self._msg_size >
MAX_TRACKER_MSG_BYTES:
+ logger.warning("Invalid message size %d from %s",
self._msg_size, self.name())
+ self.close()
+ return
+ if len(self._data) < self._msg_size:
return
+ msg = py_str(bytes(self._data[: self._msg_size]))
+ del self._data[: self._msg_size]
+ self._msg_size = 0
+ self.call_handler(json.loads(msg))
Review Comment:
The call to `json.loads(msg)` and `self.call_handler(args)` is not wrapped
in a `try...except` block. Malformed JSON or an unexpected exception in the
handler could crash the tracker or leave the connection in an inconsistent
state. Since this PR aims to improve the robustness of the tracker against
invalid or malicious inputs, it is recommended to handle these potential errors
gracefully by logging a warning and closing the connection.
##########
python/tvm/rpc/tracker.py:
##########
@@ -222,18 +224,21 @@ def on_message(self, message):
while True:
Review 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.
```suggestion
while self._sock:
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
