receiver.py 5.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200
  1. ##############################################################################
  2. #
  3. # Copyright (c) 2001, 2002 Zope Foundation and Contributors.
  4. # All Rights Reserved.
  5. #
  6. # This software is subject to the provisions of the Zope Public License,
  7. # Version 2.1 (ZPL). A copy of the ZPL should accompany this distribution.
  8. # THIS SOFTWARE IS PROVIDED "AS IS" AND ANY AND ALL EXPRESS OR IMPLIED
  9. # WARRANTIES ARE DISCLAIMED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
  10. # WARRANTIES OF TITLE, MERCHANTABILITY, AGAINST INFRINGEMENT, AND FITNESS
  11. # FOR A PARTICULAR PURPOSE.
  12. #
  13. ##############################################################################
  14. """Data Chunk Receiver
  15. """
  16. from waitress.rfc7230 import CHUNK_EXT_RE, ONLY_HEXDIG_RE
  17. from waitress.utilities import BadRequest, find_double_newline
  18. class FixedStreamReceiver:
  19. # See IStreamConsumer
  20. completed = False
  21. error = None
  22. def __init__(self, cl, buf):
  23. self.remain = cl
  24. self.buf = buf
  25. def __len__(self):
  26. return self.buf.__len__()
  27. def received(self, data):
  28. "See IStreamConsumer"
  29. rm = self.remain
  30. if rm < 1:
  31. self.completed = True # Avoid any chance of spinning
  32. return 0
  33. datalen = len(data)
  34. if rm <= datalen:
  35. self.buf.append(data[:rm])
  36. self.remain = 0
  37. self.completed = True
  38. return rm
  39. else:
  40. self.buf.append(data)
  41. self.remain -= datalen
  42. return datalen
  43. def getfile(self):
  44. return self.buf.getfile()
  45. def getbuf(self):
  46. return self.buf
  47. class ChunkedReceiver:
  48. chunk_remainder = 0
  49. validate_chunk_end = False
  50. control_line = b""
  51. chunk_end = b""
  52. all_chunks_received = False
  53. trailer = b""
  54. completed = False
  55. error = None
  56. # max_control_line = 1024
  57. # max_trailer = 65536
  58. def __init__(self, buf):
  59. self.buf = buf
  60. def __len__(self):
  61. return self.buf.__len__()
  62. def received(self, s):
  63. # Returns the number of bytes consumed.
  64. if self.completed:
  65. return 0
  66. orig_size = len(s)
  67. while s:
  68. rm = self.chunk_remainder
  69. if rm > 0:
  70. # Receive the remainder of a chunk.
  71. to_write = s[:rm]
  72. self.buf.append(to_write)
  73. written = len(to_write)
  74. s = s[written:]
  75. self.chunk_remainder -= written
  76. if self.chunk_remainder == 0:
  77. self.validate_chunk_end = True
  78. elif self.validate_chunk_end:
  79. s = self.chunk_end + s
  80. pos = s.find(b"\r\n")
  81. if pos < 0 and len(s) < 2:
  82. self.chunk_end = s
  83. s = b""
  84. else:
  85. self.chunk_end = b""
  86. if pos == 0:
  87. # Chop off the terminating CR LF from the chunk
  88. s = s[2:]
  89. else:
  90. self.error = BadRequest("Chunk not properly terminated")
  91. self.all_chunks_received = True
  92. # Always exit this loop
  93. self.validate_chunk_end = False
  94. elif not self.all_chunks_received:
  95. # Receive a control line.
  96. s = self.control_line + s
  97. pos = s.find(b"\r\n")
  98. if pos < 0:
  99. # Control line not finished.
  100. self.control_line = s
  101. s = b""
  102. else:
  103. # Control line finished.
  104. line = s[:pos]
  105. s = s[pos + 2 :]
  106. self.control_line = b""
  107. if line:
  108. # Begin a new chunk.
  109. semi = line.find(b";")
  110. if semi >= 0:
  111. extinfo = line[semi:]
  112. valid_ext_info = CHUNK_EXT_RE.match(extinfo)
  113. if not valid_ext_info:
  114. self.error = BadRequest("Invalid chunk extension")
  115. self.all_chunks_received = True
  116. break
  117. line = line[:semi]
  118. if not ONLY_HEXDIG_RE.match(line):
  119. self.error = BadRequest("Invalid chunk size")
  120. self.all_chunks_received = True
  121. break
  122. # Can not fail due to matching against the regular
  123. # expression above
  124. sz = int(line, 16) # hexadecimal
  125. if sz > 0:
  126. # Start a new chunk.
  127. self.chunk_remainder = sz
  128. else:
  129. # Finished chunks.
  130. self.all_chunks_received = True
  131. # else expect a control line.
  132. else:
  133. # Receive the trailer.
  134. trailer = self.trailer + s
  135. if trailer.startswith(b"\r\n"):
  136. # No trailer.
  137. self.completed = True
  138. return orig_size - (len(trailer) - 2)
  139. pos = find_double_newline(trailer)
  140. if pos < 0:
  141. # Trailer not finished.
  142. self.trailer = trailer
  143. s = b""
  144. else:
  145. # Finished the trailer.
  146. self.completed = True
  147. self.trailer = trailer[:pos]
  148. return orig_size - (len(trailer) - pos)
  149. return orig_size
  150. def getfile(self):
  151. return self.buf.getfile()
  152. def getbuf(self):
  153. return self.buf