| OLD | NEW |
| (Empty) |
| 1 # Copyright (c) 2010 The Chromium OS Authors. All rights reserved. | |
| 2 # Use of this source code is governed by a BSD-style license that can be | |
| 3 # found in the LICENSE file. | |
| 4 | |
| 5 """Module containing various classes pertaining to inserting a proxy in a test. | |
| 6 """ | |
| 7 | |
| 8 import os | |
| 9 import select | |
| 10 import socket | |
| 11 import SocketServer | |
| 12 import threading | |
| 13 | |
| 14 class Filter(object): | |
| 15 """Base class for data filters. | |
| 16 | |
| 17 Pass subclass of this to CrosTestProxy which will perform whatever | |
| 18 connection manipulation you prefer. | |
| 19 """ | |
| 20 | |
| 21 def setup(self): | |
| 22 """This setup method is called once per connection.""" | |
| 23 pass | |
| 24 | |
| 25 def InBound(self, data): | |
| 26 """This method is called once per packet of incoming data. | |
| 27 | |
| 28 The value returned is what is sent through the proxy. If | |
| 29 None is returned, the connection will be closed. | |
| 30 """ | |
| 31 return data | |
| 32 | |
| 33 def OutBound(self, data): | |
| 34 """This method is called once per packet of outgoing data. | |
| 35 | |
| 36 The value returned is what is sent through the proxy. If | |
| 37 None is returned, the connection will be closed. | |
| 38 """ | |
| 39 return data | |
| 40 | |
| 41 | |
| 42 class CrosTestProxy(SocketServer.ThreadingMixIn, SocketServer.TCPServer): | |
| 43 """A transparent proxy for simulating network errors""" | |
| 44 | |
| 45 class _Handler(SocketServer.BaseRequestHandler): | |
| 46 """Proxy connection handler that passes data though a filter""" | |
| 47 | |
| 48 def setup(self): | |
| 49 """Setup is called once for each connection proxied.""" | |
| 50 self.server.filter.setup() | |
| 51 | |
| 52 def handle(self): | |
| 53 """Handles each incoming connection. | |
| 54 | |
| 55 Opens a new connection to the port we are proxing to, then | |
| 56 passes each packet along in both directions after passing | |
| 57 them through the filter object passed in. | |
| 58 """ | |
| 59 # Open outgoing socket | |
| 60 s_in = self.request | |
| 61 s_out = socket.socket(socket.AF_INET, socket.SOCK_STREAM) | |
| 62 s_out.connect((self.server.address_out, self.server.port_out)) | |
| 63 | |
| 64 while True: | |
| 65 rlist, wlist, xlist = select.select([s_in, s_out], [], []) | |
| 66 | |
| 67 if s_in in rlist: | |
| 68 data = s_in.recv(1024) | |
| 69 data = self.server.filter.InBound(data) | |
| 70 if not data: break | |
| 71 try: | |
| 72 # If there is any error sending data, close both connections. | |
| 73 s_out.sendall(data) | |
| 74 except socket.error: | |
| 75 break | |
| 76 | |
| 77 if s_out in rlist: | |
| 78 data = s_out.recv(1024) | |
| 79 data = self.server.filter.OutBound(data) | |
| 80 if not data: break | |
| 81 try: | |
| 82 # If there is any error sending data, close both connections. | |
| 83 s_in.sendall(data) | |
| 84 except socket.error: | |
| 85 break | |
| 86 | |
| 87 s_in.close() | |
| 88 s_out.close() | |
| 89 | |
| 90 def __init__(self, | |
| 91 filter, | |
| 92 port_in=8081, | |
| 93 address_out='127.0.0.1', port_out=8080): | |
| 94 """Configures the proxy object. | |
| 95 | |
| 96 Args: | |
| 97 filter: An instance of a subclass of Filter. | |
| 98 port_in: Port on which to listen for incoming connections. | |
| 99 address_out: Address to which outgoing connections will go. | |
| 100 address_port: Port to which outgoing connections will go. | |
| 101 """ | |
| 102 self.port_in = port_in | |
| 103 self.address_out = address_out | |
| 104 self.port_out = port_out | |
| 105 self.filter = filter | |
| 106 | |
| 107 try: | |
| 108 SocketServer.TCPServer.__init__(self, | |
| 109 ('', port_in), | |
| 110 self._Handler) | |
| 111 except socket.error: | |
| 112 os.system('sudo netstat -l --tcp -n -p') | |
| 113 raise | |
| 114 | |
| 115 def serve_forever_in_thread(self): | |
| 116 """Helper method to start the server in a new background thread.""" | |
| 117 server_thread = threading.Thread(target=self.serve_forever) | |
| 118 server_thread.setDaemon(True) | |
| 119 server_thread.start() | |
| 120 | |
| 121 return server_thread | |
| OLD | NEW |