simplify eof detection

This commit is contained in:
Maximilian Hils 2016-10-05 17:45:42 -07:00
parent 8e7ec6117a
commit b422ba5328
2 changed files with 5 additions and 16 deletions

View File

@ -172,6 +172,8 @@ def load(file_handle):
# Read the length prefix one char at a time.
# Note that the netstring spec explicitly forbids padding zeros.
c = file_handle.read(1)
if c == b"": # we want to detect this special case.
raise ValueError("not a tnetstring: empty file")
data_length = b""
while c.isdigit():
data_length += c

View File

@ -26,16 +26,6 @@ class FlowReader:
"""
Yields Flow objects from the dump.
"""
# There is a weird mingw bug that breaks .tell() when reading from stdin.
try:
self.fo.tell()
except IOError: # pragma: no cover
can_tell = False
else:
can_tell = True
off = 0
try:
while True:
data = tnetstring.load(self.fo)
@ -43,15 +33,12 @@ class FlowReader:
data = io_compat.migrate_flow(data)
except ValueError as e:
raise exceptions.FlowReadException(str(e))
if can_tell:
off = self.fo.tell()
if data["type"] not in models.FLOW_TYPES:
raise exceptions.FlowReadException("Unknown flow type: {}".format(data["type"]))
yield models.FLOW_TYPES[data["type"]].from_state(data)
except ValueError:
# Error is due to EOF
if can_tell and self.fo.tell() == off and self.fo.read() == b'':
return
except ValueError as e:
if str(e) == "not a tnetstring: empty file":
return # Error is due to EOF
raise exceptions.FlowReadException("Invalid data format.")