Traffic.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446
  1. #!/usr/bin/env python
  2. #
  3. # Copyright 2013 The Tor Project
  4. #
  5. # You may do anything with this work that copyright law would normally
  6. # restrict, so long as you retain the above notice(s) and this license
  7. # in all redistributed copies and derived works. There is no warranty.
  8. # Do select/read/write for binding to a port, connecting to it and
  9. # write, read what's written and verify it. You can connect over a
  10. # SOCKS proxy (like Tor).
  11. #
  12. # You can create a TrafficTester and give it an IP address/host and
  13. # port to bind to. If a Source is created and added to the
  14. # TrafficTester, it will connect to the address/port it was given at
  15. # instantiation and send its data. A Source can be configured to
  16. # connect over a SOCKS proxy. When everything is set up, you can
  17. # invoke TrafficTester.run() to start running. The TrafficTester will
  18. # accept the incoming connection and read from it, verifying the data.
  19. #
  20. # For example code, see main() below.
  21. from __future__ import print_function
  22. import sys
  23. import socket
  24. import select
  25. import struct
  26. import errno
  27. import time
  28. import os
  29. from chutney.Debug import debug_flag, debug
  30. def socks_cmd(addr_port):
  31. """
  32. Return a SOCKS command for connecting to addr_port.
  33. SOCKSv4: https://en.wikipedia.org/wiki/SOCKS#Protocol
  34. SOCKSv5: RFC1928, RFC1929
  35. """
  36. ver = 4 # Only SOCKSv4 for now.
  37. cmd = 1 # Stream connection.
  38. user = b'\x00'
  39. dnsname = ''
  40. host, port = addr_port
  41. try:
  42. addr = socket.inet_aton(host)
  43. except socket.error:
  44. addr = b'\x00\x00\x00\x01'
  45. dnsname = '%s\x00' % host
  46. debug("Socks 4a request to %s:%d" % (host, port))
  47. if type(dnsname) != type(b""):
  48. dnsname = dnsname.encode("ascii")
  49. return struct.pack('!BBH', ver, cmd, port) + addr + user + dnsname
  50. class TestSuite(object):
  51. """Keep a tab on how many tests are pending, how many have failed
  52. and how many have succeeded."""
  53. def __init__(self):
  54. self.not_done = 0
  55. self.successes = 0
  56. self.failures = 0
  57. def add(self):
  58. self.not_done += 1
  59. def success(self):
  60. self.not_done -= 1
  61. self.successes += 1
  62. def failure(self):
  63. self.not_done -= 1
  64. self.failures += 1
  65. def failure_count(self):
  66. return self.failures
  67. def all_done(self):
  68. return self.not_done == 0
  69. def status(self):
  70. return('%d/%d/%d' % (self.not_done, self.successes, self.failures))
  71. class Peer(object):
  72. "Base class for Listener, Source and Sink."
  73. LISTENER = 1
  74. SOURCE = 2
  75. SINK = 3
  76. def __init__(self, ptype, tt, s=None):
  77. self.type = ptype
  78. self.tt = tt # TrafficTester
  79. if s is not None:
  80. self.s = s
  81. else:
  82. self.s = socket.socket()
  83. self.s.setblocking(False)
  84. def fd(self):
  85. return self.s.fileno()
  86. def is_source(self):
  87. return self.type == self.SOURCE
  88. def is_sink(self):
  89. return self.type == self.SINK
  90. class Listener(Peer):
  91. "A TCP listener, binding, listening and accepting new connections."
  92. def __init__(self, tt, endpoint):
  93. super(Listener, self).__init__(Peer.LISTENER, tt)
  94. self.s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  95. self.s.bind(endpoint)
  96. self.s.listen(0)
  97. def accept(self):
  98. newsock, endpoint = self.s.accept()
  99. debug("new client from %s:%s (fd=%d)" %
  100. (endpoint[0], endpoint[1], newsock.fileno()))
  101. self.tt.add(Sink(self.tt, newsock))
  102. class Sink(Peer):
  103. "A data sink, reading from its peer and verifying the data."
  104. def __init__(self, tt, s):
  105. super(Sink, self).__init__(Peer.SINK, tt, s)
  106. self.inbuf = b''
  107. self.repetitions = self.tt.repetitions
  108. def on_readable(self):
  109. """Invoked when the socket becomes readable.
  110. Return 0 on finished, successful verification.
  111. -1 on failed verification
  112. >0 if more data needs to be read
  113. """
  114. return self.verify(self.tt.data)
  115. def verify(self, data):
  116. # shortcut read when we don't ever expect any data
  117. if self.repetitions == 0 or len(self.tt.data) == 0:
  118. debug("no verification required - no data")
  119. return 0
  120. inp = self.s.recv(len(data) - len(self.inbuf))
  121. debug("Verify: received %d bytes"% len(inp))
  122. if len(inp) == 0:
  123. debug("EOF on fd %s" % self.fd())
  124. return -1
  125. self.inbuf += inp
  126. debug("successfully received (bytes=%d)" % len(self.inbuf))
  127. while len(self.inbuf) >= len(data):
  128. assert(len(self.inbuf) <= len(data) or self.repetitions > 1)
  129. if self.inbuf[:len(data)] != data:
  130. debug("receive comparison failed (bytes=%d)" % len(data))
  131. return -1 # Failed verification.
  132. # if we're not debugging, print a dot every dot_repetitions reps
  133. elif (not debug_flag and self.tt.dot_repetitions > 0 and
  134. self.repetitions % self.tt.dot_repetitions == 0):
  135. sys.stdout.write('.')
  136. sys.stdout.flush()
  137. # repeatedly check data against self.inbuf if required
  138. debug("receive comparison success (bytes=%d)" % len(data))
  139. self.inbuf = self.inbuf[len(data):]
  140. debug("receive leftover bytes (bytes=%d)" % len(self.inbuf))
  141. self.repetitions -= 1
  142. debug("receive remaining repetitions (reps=%d)" % self.repetitions)
  143. if self.repetitions == 0 and len(self.inbuf) == 0:
  144. debug("successful verification")
  145. # calculate the actual length of data remaining, including reps
  146. debug("receive remaining bytes (bytes=%d)"
  147. % (self.repetitions*len(data) - len(self.inbuf)))
  148. return self.repetitions*len(data) - len(self.inbuf)
  149. class Source(Peer):
  150. """A data source, connecting to a TCP server, optionally over a
  151. SOCKS proxy, sending data."""
  152. NOT_CONNECTED = 0
  153. CONNECTING = 1
  154. CONNECTING_THROUGH_PROXY = 2
  155. CONNECTED = 5
  156. def __init__(self, tt, server, buf, proxy=None, repetitions=1):
  157. super(Source, self).__init__(Peer.SOURCE, tt)
  158. self.state = self.NOT_CONNECTED
  159. self.data = buf
  160. self.outbuf = b''
  161. self.inbuf = b''
  162. self.proxy = proxy
  163. self.repetitions = repetitions
  164. self._sent_no_bytes = 0
  165. # sanity checks
  166. if len(self.data) == 0:
  167. self.repetitions = 0
  168. if self.repetitions == 0:
  169. self.data = {}
  170. self.connect(server)
  171. def connect(self, endpoint):
  172. self.dest = endpoint
  173. self.state = self.CONNECTING
  174. dest = self.proxy or self.dest
  175. try:
  176. debug("socket %d connecting to %r..."%(self.fd(),dest))
  177. self.s.connect(dest)
  178. except socket.error as e:
  179. if e.errno != errno.EINPROGRESS:
  180. raise
  181. def on_readable(self):
  182. """Invoked when the socket becomes readable.
  183. Return -1 on failure
  184. >0 if more data needs to be read or written
  185. """
  186. if self.state == self.CONNECTING_THROUGH_PROXY:
  187. inp = self.s.recv(8 - len(self.inbuf))
  188. debug("-- connecting through proxy, got %d bytes"%len(inp))
  189. if len(inp) == 0:
  190. debug("EOF on fd %d"%self.fd())
  191. return -1
  192. self.inbuf += inp
  193. if len(self.inbuf) == 8:
  194. if self.inbuf[:2] == b'\x00\x5a':
  195. debug("proxy handshake successful (fd=%d)" % self.fd())
  196. self.state = self.CONNECTED
  197. self.inbuf = b''
  198. debug("successfully connected (fd=%d)" % self.fd())
  199. # if we have no reps or no data, skip sending actual data
  200. if self.want_to_write():
  201. return 1 # Keep us around for writing.
  202. else:
  203. # shortcut write when we don't ever expect any data
  204. debug("no connection required - no data")
  205. return 0
  206. else:
  207. debug("proxy handshake failed (0x%x)! (fd=%d)" %
  208. (ord(self.inbuf[1]), self.fd()))
  209. self.state = self.NOT_CONNECTED
  210. return -1
  211. assert(8 - len(self.inbuf) > 0)
  212. return 8 - len(self.inbuf)
  213. return self.want_to_write() # Keep us around for writing if needed
  214. def want_to_write(self):
  215. if self.state == self.CONNECTING:
  216. return True
  217. if len(self.outbuf) > 0:
  218. return True
  219. if (self.state == self.CONNECTED and
  220. self.repetitions > 0 and
  221. len(self.data) > 0):
  222. return True
  223. return False
  224. def on_writable(self):
  225. """Invoked when the socket becomes writable.
  226. Return 0 when done writing
  227. -1 on failure (like connection refused)
  228. >0 if more data needs to be written
  229. """
  230. if self.state == self.CONNECTING:
  231. if self.proxy is None:
  232. self.state = self.CONNECTED
  233. debug("successfully connected (fd=%d)" % self.fd())
  234. else:
  235. self.state = self.CONNECTING_THROUGH_PROXY
  236. self.outbuf = socks_cmd(self.dest)
  237. # we write socks_cmd() to the proxy, then read the response
  238. # if we get the correct response, we're CONNECTED
  239. if self.state == self.CONNECTED:
  240. # repeat self.data into self.outbuf if required
  241. if (len(self.outbuf) < len(self.data) and self.repetitions > 0):
  242. self.outbuf += self.data
  243. self.repetitions -= 1
  244. debug("adding more data to send (bytes=%d)" % len(self.data))
  245. debug("now have data to send (bytes=%d)" % len(self.outbuf))
  246. debug("send repetitions remaining (reps=%d)"
  247. % self.repetitions)
  248. try:
  249. n = self.s.send(self.outbuf)
  250. except socket.error as e:
  251. if e.errno == errno.ECONNREFUSED:
  252. debug("connection refused (fd=%d)" % self.fd())
  253. return -1
  254. raise
  255. # sometimes, this debug statement prints 0
  256. # it should print length of the data sent
  257. # but the code works as long as this doesn't keep on happening
  258. if n > 0:
  259. debug("successfully sent (bytes=%d)" % n)
  260. self._sent_no_bytes = 0
  261. else:
  262. debug("BUG: sent no bytes (out of %d; state is %s)"% (len(self.outbuf), self.state))
  263. self._sent_no_bytes += 1
  264. # We can't retry too fast, otherwise clients burn all their HSDirs
  265. if self._sent_no_bytes >= 2:
  266. print("Sent no data %d times. Stalled." %
  267. (self._sent_no_bytes))
  268. return -1
  269. time.sleep(5)
  270. self.outbuf = self.outbuf[n:]
  271. if self.state == self.CONNECTING_THROUGH_PROXY:
  272. return 1 # Keep us around.
  273. debug("bytes remaining on outbuf (bytes=%d)" % len(self.outbuf))
  274. # calculate the actual length of data remaining, including reps
  275. # When 0, we're being removed.
  276. debug("bytes remaining overall (bytes=%d)"
  277. % (self.repetitions*len(self.data) + len(self.outbuf)))
  278. return self.repetitions*len(self.data) + len(self.outbuf)
  279. class TrafficTester():
  280. """
  281. Hang on select.select() and dispatch to Sources and Sinks.
  282. Time out after self.timeout seconds.
  283. Keep track of successful and failed data verification using a
  284. TestSuite.
  285. Return True if all tests succeed, else False.
  286. """
  287. def __init__(self,
  288. endpoint,
  289. data={},
  290. timeout=3,
  291. repetitions=1,
  292. dot_repetitions=0):
  293. self.listener = Listener(self, endpoint)
  294. self.pending_close = []
  295. self.timeout = timeout
  296. self.tests = TestSuite()
  297. self.data = data
  298. self.repetitions = repetitions
  299. # sanity checks
  300. if len(self.data) == 0:
  301. self.repetitions = 0
  302. if self.repetitions == 0:
  303. self.data = {}
  304. self.dot_repetitions = dot_repetitions
  305. debug("listener fd=%d" % self.listener.fd())
  306. self.peers = {} # fd:Peer
  307. def sinks(self):
  308. return self.get_by_ptype(Peer.SINK)
  309. def sources(self):
  310. return self.get_by_ptype(Peer.SOURCE)
  311. def get_by_ptype(self, ptype):
  312. return list(filter(lambda p: p.type == ptype, self.peers.values()))
  313. def add(self, peer):
  314. self.peers[peer.fd()] = peer
  315. if peer.is_source():
  316. self.tests.add()
  317. def remove(self, peer):
  318. self.peers.pop(peer.fd())
  319. self.pending_close.append(peer.s)
  320. def run(self):
  321. while not self.tests.all_done() and self.timeout > 0:
  322. rset = [self.listener.fd()] + list(self.peers)
  323. wset = [p.fd() for p in
  324. filter(lambda x: x.want_to_write(), self.sources())]
  325. # debug("rset %s wset %s" % (rset, wset))
  326. sets = select.select(rset, wset, [], 1)
  327. if all(len(s) == 0 for s in sets):
  328. debug("Decrementing timeout.")
  329. self.timeout -= 1
  330. continue
  331. for fd in sets[0]: # readable fd's
  332. if fd == self.listener.fd():
  333. self.listener.accept()
  334. continue
  335. p = self.peers[fd]
  336. n = p.on_readable()
  337. debug("On read, fd %d for %s said %d"%(fd, p, n))
  338. if n > 0:
  339. # debug("need %d more octets from fd %d" % (n, fd))
  340. pass
  341. elif n == 0: # Success.
  342. self.tests.success()
  343. self.remove(p)
  344. else: # Failure.
  345. debug("Got a failure reading fd %d for %s" % (fd,p))
  346. self.tests.failure()
  347. if p.is_sink():
  348. print("verification failed!")
  349. self.remove(p)
  350. for fd in sets[1]: # writable fd's
  351. p = self.peers.get(fd)
  352. if p is not None: # Might have been removed above.
  353. n = p.on_writable()
  354. debug("On write, fd %d said %d"%(fd, n))
  355. if n == 0:
  356. self.remove(p)
  357. elif n < 0:
  358. debug("Got a failure writing fd %d for %s" % (fd,p))
  359. self.tests.failure()
  360. self.remove(p)
  361. for fd in self.peers:
  362. peer = self.peers[fd]
  363. debug("peer fd=%d never pending close, never read or wrote" % fd)
  364. self.pending_close.append(peer.s)
  365. self.listener.s.close()
  366. for s in self.pending_close:
  367. s.close()
  368. if not debug_flag:
  369. sys.stdout.write('\n')
  370. sys.stdout.flush()
  371. debug("Done with run(); all_done == %s and failure_count == %s"
  372. %(self.tests.all_done(), self.tests.failure_count()))
  373. return self.tests.all_done() and self.tests.failure_count() == 0
  374. def main():
  375. """Test the TrafficTester by sending and receiving some data."""
  376. DATA = b"a foo is a bar" * 1000
  377. bind_to = ('localhost', int(sys.argv[1]))
  378. tt = TrafficTester(bind_to, DATA)
  379. # Don't use a proxy for self-testing, so that we avoid tor entirely
  380. tt.add(Source(tt, bind_to, DATA))
  381. success = tt.run()
  382. if success:
  383. return 0
  384. return 255
  385. if __name__ == '__main__':
  386. sys.exit(main())