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