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 |