gemini-code-assist[bot] commented on code in PR #19586:
URL: https://github.com/apache/tvm/pull/19586#discussion_r3258351973
##########
python/tvm/rpc/tracker.py:
##########
@@ -224,11 +230,20 @@ def on_message(self, message):
if self._msg_size == 0:
if len(self._data) >= 4:
self._msg_size = struct.unpack("<i", self._data[:4])[0]
+ if self._msg_size <= 0 or self._msg_size >
MAX_TRACKER_MSG_BYTES:
+ logger.warning(
+ "Invalid msg_size %d from %s; closing connection",
+ self._msg_size,
+ self.name(),
+ )
+ self.close()
+ return
+ del self._data[:4]
else:
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]
+ if self._msg_size != 0 and len(self._data) >= self._msg_size:
+ msg = py_str(bytes(self._data[: self._msg_size]))
+ del self._data[: self._msg_size]
self._msg_size = 0
# pylint: disable=broad-except
self.call_handler(json.loads(msg))
Review Comment:
The `pylint: disable=broad-except` comment suggests that an exception
handler was intended here, but the `try...except` block is missing. Malformed
JSON input or errors in `call_handler` could cause the tracker to raise an
unhandled exception, potentially leading to an ungraceful connection
termination or affecting the event loop. Given the focus on robustness in this
PR, it would be better to wrap this call in a `try...except` block and close
the connection on failure.
```suggestion
try:
self.call_handler(json.loads(msg))
except Exception: # pylint: disable=broad-except
logger.warning("Error in handling message from %s",
self.name())
self.close()
return
```
--
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]
