Traffic.py 14 KB

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