test_rawudp.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327
  1. # Copyright (c) Twisted Matrix Laboratories.
  2. # See LICENSE for details.
  3. #
  4. from twisted.trial import unittest
  5. from twisted.internet import protocol
  6. from twisted.pair import rawudp
  7. class MyProtocol(protocol.DatagramProtocol):
  8. def __init__(self, expecting):
  9. self.expecting = list(expecting)
  10. def datagramReceived(self, data, peer):
  11. (host, port) = peer
  12. assert self.expecting, 'Got a packet when not expecting anymore.'
  13. expectData, expectHost, expectPort = self.expecting.pop(0)
  14. assert expectData == data, "Expected data %r, got %r" % (expectData, data)
  15. assert expectHost == host, "Expected host %r, got %r" % (expectHost, host)
  16. assert expectPort == port, "Expected port %d=0x%04x, got %d=0x%04x" % (expectPort, expectPort, port, port)
  17. class RawUDPTests(unittest.TestCase):
  18. def testPacketParsing(self):
  19. proto = rawudp.RawUDPProtocol()
  20. p1 = MyProtocol([
  21. (b'foobar', b'testHost', 0x43A2),
  22. ])
  23. proto.addProto(0xF00F, p1)
  24. proto.datagramReceived(b"\x43\xA2" #source
  25. b"\xf0\x0f" #dest
  26. b"\x00\x06" #len
  27. b"\xDE\xAD" #check
  28. b"foobar",
  29. partial=0,
  30. dest=b'dummy',
  31. source=b'testHost',
  32. protocol=b'dummy',
  33. version=b'dummy',
  34. ihl=b'dummy',
  35. tos=b'dummy',
  36. tot_len=b'dummy',
  37. fragment_id=b'dummy',
  38. fragment_offset=b'dummy',
  39. dont_fragment=b'dummy',
  40. more_fragments=b'dummy',
  41. ttl=b'dummy',
  42. )
  43. assert not p1.expecting, \
  44. 'Should not expect any more packets, but still want %r' % p1.expecting
  45. def testMultiplePackets(self):
  46. proto = rawudp.RawUDPProtocol()
  47. p1 = MyProtocol([
  48. (b'foobar', b'testHost', 0x43A2),
  49. (b'quux', b'otherHost', 0x33FE),
  50. ])
  51. proto.addProto(0xF00F, p1)
  52. proto.datagramReceived(b"\x43\xA2" #source
  53. b"\xf0\x0f" #dest
  54. b"\x00\x06" #len
  55. b"\xDE\xAD" #check
  56. b"foobar",
  57. partial=0,
  58. dest=b'dummy',
  59. source=b'testHost',
  60. protocol=b'dummy',
  61. version=b'dummy',
  62. ihl=b'dummy',
  63. tos=b'dummy',
  64. tot_len=b'dummy',
  65. fragment_id=b'dummy',
  66. fragment_offset=b'dummy',
  67. dont_fragment=b'dummy',
  68. more_fragments=b'dummy',
  69. ttl=b'dummy',
  70. )
  71. proto.datagramReceived(b"\x33\xFE" #source
  72. b"\xf0\x0f" #dest
  73. b"\x00\x05" #len
  74. b"\xDE\xAD" #check
  75. b"quux",
  76. partial=0,
  77. dest=b'dummy',
  78. source=b'otherHost',
  79. protocol=b'dummy',
  80. version=b'dummy',
  81. ihl=b'dummy',
  82. tos=b'dummy',
  83. tot_len=b'dummy',
  84. fragment_id=b'dummy',
  85. fragment_offset=b'dummy',
  86. dont_fragment=b'dummy',
  87. more_fragments=b'dummy',
  88. ttl=b'dummy',
  89. )
  90. assert not p1.expecting, \
  91. 'Should not expect any more packets, but still want %r' % p1.expecting
  92. def testMultipleSameProtos(self):
  93. proto = rawudp.RawUDPProtocol()
  94. p1 = MyProtocol([
  95. (b'foobar', b'testHost', 0x43A2),
  96. ])
  97. p2 = MyProtocol([
  98. (b'foobar', b'testHost', 0x43A2),
  99. ])
  100. proto.addProto(0xF00F, p1)
  101. proto.addProto(0xF00F, p2)
  102. proto.datagramReceived(b"\x43\xA2" #source
  103. b"\xf0\x0f" #dest
  104. b"\x00\x06" #len
  105. b"\xDE\xAD" #check
  106. b"foobar",
  107. partial=0,
  108. dest=b'dummy',
  109. source=b'testHost',
  110. protocol=b'dummy',
  111. version=b'dummy',
  112. ihl=b'dummy',
  113. tos=b'dummy',
  114. tot_len=b'dummy',
  115. fragment_id=b'dummy',
  116. fragment_offset=b'dummy',
  117. dont_fragment=b'dummy',
  118. more_fragments=b'dummy',
  119. ttl=b'dummy',
  120. )
  121. assert not p1.expecting, \
  122. 'Should not expect any more packets, but still want %r' % p1.expecting
  123. assert not p2.expecting, \
  124. 'Should not expect any more packets, but still want %r' % p2.expecting
  125. def testWrongProtoNotSeen(self):
  126. proto = rawudp.RawUDPProtocol()
  127. p1 = MyProtocol([])
  128. proto.addProto(1, p1)
  129. proto.datagramReceived(b"\x43\xA2" #source
  130. b"\xf0\x0f" #dest
  131. b"\x00\x06" #len
  132. b"\xDE\xAD" #check
  133. b"foobar",
  134. partial=0,
  135. dest=b'dummy',
  136. source=b'testHost',
  137. protocol=b'dummy',
  138. version=b'dummy',
  139. ihl=b'dummy',
  140. tos=b'dummy',
  141. tot_len=b'dummy',
  142. fragment_id=b'dummy',
  143. fragment_offset=b'dummy',
  144. dont_fragment=b'dummy',
  145. more_fragments=b'dummy',
  146. ttl=b'dummy',
  147. )
  148. def testDemuxing(self):
  149. proto = rawudp.RawUDPProtocol()
  150. p1 = MyProtocol([
  151. (b'foobar', b'testHost', 0x43A2),
  152. (b'quux', b'otherHost', 0x33FE),
  153. ])
  154. proto.addProto(0xF00F, p1)
  155. p2 = MyProtocol([
  156. (b'quux', b'otherHost', 0xA401),
  157. (b'foobar', b'testHost', 0xA302),
  158. ])
  159. proto.addProto(0xB050, p2)
  160. proto.datagramReceived(b"\xA4\x01" #source
  161. b"\xB0\x50" #dest
  162. b"\x00\x05" #len
  163. b"\xDE\xAD" #check
  164. b"quux",
  165. partial=0,
  166. dest=b'dummy',
  167. source=b'otherHost',
  168. protocol=b'dummy',
  169. version=b'dummy',
  170. ihl=b'dummy',
  171. tos=b'dummy',
  172. tot_len=b'dummy',
  173. fragment_id=b'dummy',
  174. fragment_offset=b'dummy',
  175. dont_fragment=b'dummy',
  176. more_fragments=b'dummy',
  177. ttl=b'dummy',
  178. )
  179. proto.datagramReceived(b"\x43\xA2" #source
  180. b"\xf0\x0f" #dest
  181. b"\x00\x06" #len
  182. b"\xDE\xAD" #check
  183. b"foobar",
  184. partial=0,
  185. dest=b'dummy',
  186. source=b'testHost',
  187. protocol=b'dummy',
  188. version=b'dummy',
  189. ihl=b'dummy',
  190. tos=b'dummy',
  191. tot_len=b'dummy',
  192. fragment_id=b'dummy',
  193. fragment_offset=b'dummy',
  194. dont_fragment=b'dummy',
  195. more_fragments=b'dummy',
  196. ttl=b'dummy',
  197. )
  198. proto.datagramReceived(b"\x33\xFE" #source
  199. b"\xf0\x0f" #dest
  200. b"\x00\x05" #len
  201. b"\xDE\xAD" #check
  202. b"quux",
  203. partial=0,
  204. dest=b'dummy',
  205. source=b'otherHost',
  206. protocol=b'dummy',
  207. version=b'dummy',
  208. ihl=b'dummy',
  209. tos=b'dummy',
  210. tot_len=b'dummy',
  211. fragment_id=b'dummy',
  212. fragment_offset=b'dummy',
  213. dont_fragment=b'dummy',
  214. more_fragments=b'dummy',
  215. ttl=b'dummy',
  216. )
  217. proto.datagramReceived(b"\xA3\x02" #source
  218. b"\xB0\x50" #dest
  219. b"\x00\x06" #len
  220. b"\xDE\xAD" #check
  221. b"foobar",
  222. partial=0,
  223. dest=b'dummy',
  224. source=b'testHost',
  225. protocol=b'dummy',
  226. version=b'dummy',
  227. ihl=b'dummy',
  228. tos=b'dummy',
  229. tot_len=b'dummy',
  230. fragment_id=b'dummy',
  231. fragment_offset=b'dummy',
  232. dont_fragment=b'dummy',
  233. more_fragments=b'dummy',
  234. ttl=b'dummy',
  235. )
  236. assert not p1.expecting, \
  237. 'Should not expect any more packets, but still want %r' % p1.expecting
  238. assert not p2.expecting, \
  239. 'Should not expect any more packets, but still want %r' % p2.expecting
  240. def testAddingBadProtos_WrongLevel(self):
  241. """Adding a wrong level protocol raises an exception."""
  242. e = rawudp.RawUDPProtocol()
  243. try:
  244. e.addProto(42, "silliness")
  245. except TypeError as e:
  246. if e.args == ('Added protocol must be an instance of DatagramProtocol',):
  247. pass
  248. else:
  249. raise
  250. else:
  251. raise AssertionError('addProto must raise an exception for bad protocols')
  252. def testAddingBadProtos_TooSmall(self):
  253. """Adding a protocol with a negative number raises an exception."""
  254. e = rawudp.RawUDPProtocol()
  255. try:
  256. e.addProto(-1, protocol.DatagramProtocol())
  257. except TypeError as e:
  258. if e.args == ('Added protocol must be positive or zero',):
  259. pass
  260. else:
  261. raise
  262. else:
  263. raise AssertionError('addProto must raise an exception for bad protocols')
  264. def testAddingBadProtos_TooBig(self):
  265. """Adding a protocol with a number >=2**16 raises an exception."""
  266. e = rawudp.RawUDPProtocol()
  267. try:
  268. e.addProto(2**16, protocol.DatagramProtocol())
  269. except TypeError as e:
  270. if e.args == ('Added protocol must fit in 16 bits',):
  271. pass
  272. else:
  273. raise
  274. else:
  275. raise AssertionError('addProto must raise an exception for bad protocols')
  276. def testAddingBadProtos_TooBig2(self):
  277. """Adding a protocol with a number >=2**16 raises an exception."""
  278. e = rawudp.RawUDPProtocol()
  279. try:
  280. e.addProto(2**16+1, protocol.DatagramProtocol())
  281. except TypeError as e:
  282. if e.args == ('Added protocol must fit in 16 bits',):
  283. pass
  284. else:
  285. raise
  286. else:
  287. raise AssertionError('addProto must raise an exception for bad protocols')