bl4cksku11 commented on code in PR #19586:
URL: https://github.com/apache/tvm/pull/19586#discussion_r3258386117
##########
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:
Fair. Wrapped json.loads + call_handler in try/except and added
exc_info=True so the traceback gets logged for diagnostics. Pushed as 8f6875a.
--
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]
