##// END OF EJS Templates
py3: add b'' prefixes to config options in test/badserverext.py...
Pulkit Goyal -
r36413:fb7897e5 default
parent child Browse files
Show More
@@ -1,290 +1,290
1 1 # badserverext.py - Extension making servers behave badly
2 2 #
3 3 # Copyright 2017 Gregory Szorc <gregory.szorc@gmail.com>
4 4 #
5 5 # This software may be used and distributed according to the terms of the
6 6 # GNU General Public License version 2 or any later version.
7 7
8 8 # no-check-code
9 9
10 10 """Extension to make servers behave badly.
11 11
12 12 This extension is useful for testing Mercurial behavior when various network
13 13 events occur.
14 14
15 15 Various config options in the [badserver] section influence behavior:
16 16
17 17 closebeforeaccept
18 18 If true, close() the server socket when a new connection arrives before
19 19 accept() is called. The server will then exit.
20 20
21 21 closeafteraccept
22 22 If true, the server will close() the client socket immediately after
23 23 accept().
24 24
25 25 closeafterrecvbytes
26 26 If defined, close the client socket after receiving this many bytes.
27 27
28 28 closeaftersendbytes
29 29 If defined, close the client socket after sending this many bytes.
30 30 """
31 31
32 32 from __future__ import absolute_import
33 33
34 34 import socket
35 35
36 36 from mercurial import(
37 37 registrar,
38 38 )
39 39
40 40 from mercurial.hgweb import (
41 41 server,
42 42 )
43 43
44 44 configtable = {}
45 45 configitem = registrar.configitem(configtable)
46 46
47 configitem('badserver', 'closeafteraccept',
47 configitem(b'badserver', b'closeafteraccept',
48 48 default=False,
49 49 )
50 configitem('badserver', 'closeafterrecvbytes',
50 configitem(b'badserver', b'closeafterrecvbytes',
51 51 default=0,
52 52 )
53 configitem('badserver', 'closeaftersendbytes',
53 configitem(b'badserver', b'closeaftersendbytes',
54 54 default=0,
55 55 )
56 configitem('badserver', 'closebeforeaccept',
56 configitem(b'badserver', b'closebeforeaccept',
57 57 default=False,
58 58 )
59 59
60 60 # We can't adjust __class__ on a socket instance. So we define a proxy type.
61 61 class socketproxy(object):
62 62 __slots__ = (
63 63 '_orig',
64 64 '_logfp',
65 65 '_closeafterrecvbytes',
66 66 '_closeaftersendbytes',
67 67 )
68 68
69 69 def __init__(self, obj, logfp, closeafterrecvbytes=0,
70 70 closeaftersendbytes=0):
71 71 object.__setattr__(self, '_orig', obj)
72 72 object.__setattr__(self, '_logfp', logfp)
73 73 object.__setattr__(self, '_closeafterrecvbytes', closeafterrecvbytes)
74 74 object.__setattr__(self, '_closeaftersendbytes', closeaftersendbytes)
75 75
76 76 def __getattribute__(self, name):
77 77 if name in ('makefile',):
78 78 return object.__getattribute__(self, name)
79 79
80 80 return getattr(object.__getattribute__(self, '_orig'), name)
81 81
82 82 def __delattr__(self, name):
83 83 delattr(object.__getattribute__(self, '_orig'), name)
84 84
85 85 def __setattr__(self, name, value):
86 86 setattr(object.__getattribute__(self, '_orig'), name, value)
87 87
88 88 def makefile(self, mode, bufsize):
89 89 f = object.__getattribute__(self, '_orig').makefile(mode, bufsize)
90 90
91 91 logfp = object.__getattribute__(self, '_logfp')
92 92 closeafterrecvbytes = object.__getattribute__(self,
93 93 '_closeafterrecvbytes')
94 94 closeaftersendbytes = object.__getattribute__(self,
95 95 '_closeaftersendbytes')
96 96
97 97 return fileobjectproxy(f, logfp,
98 98 closeafterrecvbytes=closeafterrecvbytes,
99 99 closeaftersendbytes=closeaftersendbytes)
100 100
101 101 # We can't adjust __class__ on socket._fileobject, so define a proxy.
102 102 class fileobjectproxy(object):
103 103 __slots__ = (
104 104 '_orig',
105 105 '_logfp',
106 106 '_closeafterrecvbytes',
107 107 '_closeaftersendbytes',
108 108 )
109 109
110 110 def __init__(self, obj, logfp, closeafterrecvbytes=0,
111 111 closeaftersendbytes=0):
112 112 object.__setattr__(self, '_orig', obj)
113 113 object.__setattr__(self, '_logfp', logfp)
114 114 object.__setattr__(self, '_closeafterrecvbytes', closeafterrecvbytes)
115 115 object.__setattr__(self, '_closeaftersendbytes', closeaftersendbytes)
116 116
117 117 def __getattribute__(self, name):
118 118 if name in ('read', 'readline', 'write', '_writelog'):
119 119 return object.__getattribute__(self, name)
120 120
121 121 return getattr(object.__getattribute__(self, '_orig'), name)
122 122
123 123 def __delattr__(self, name):
124 124 delattr(object.__getattribute__(self, '_orig'), name)
125 125
126 126 def __setattr__(self, name, value):
127 127 setattr(object.__getattribute__(self, '_orig'), name, value)
128 128
129 129 def _writelog(self, msg):
130 130 msg = msg.replace('\r', '\\r').replace('\n', '\\n')
131 131
132 132 object.__getattribute__(self, '_logfp').write(msg)
133 133 object.__getattribute__(self, '_logfp').write('\n')
134 134 object.__getattribute__(self, '_logfp').flush()
135 135
136 136 def read(self, size=-1):
137 137 remaining = object.__getattribute__(self, '_closeafterrecvbytes')
138 138
139 139 # No read limit. Call original function.
140 140 if not remaining:
141 141 result = object.__getattribute__(self, '_orig').read(size)
142 142 self._writelog('read(%d) -> (%d) (%s) %s' % (size,
143 143 len(result),
144 144 result))
145 145 return result
146 146
147 147 origsize = size
148 148
149 149 if size < 0:
150 150 size = remaining
151 151 else:
152 152 size = min(remaining, size)
153 153
154 154 result = object.__getattribute__(self, '_orig').read(size)
155 155 remaining -= len(result)
156 156
157 157 self._writelog('read(%d from %d) -> (%d) %s' % (
158 158 size, origsize, len(result), result))
159 159
160 160 object.__setattr__(self, '_closeafterrecvbytes', remaining)
161 161
162 162 if remaining <= 0:
163 163 self._writelog('read limit reached, closing socket')
164 164 self._sock.close()
165 165 # This is the easiest way to abort the current request.
166 166 raise Exception('connection closed after receiving N bytes')
167 167
168 168 return result
169 169
170 170 def readline(self, size=-1):
171 171 remaining = object.__getattribute__(self, '_closeafterrecvbytes')
172 172
173 173 # No read limit. Call original function.
174 174 if not remaining:
175 175 result = object.__getattribute__(self, '_orig').readline(size)
176 176 self._writelog('readline(%d) -> (%d) %s' % (
177 177 size, len(result), result))
178 178 return result
179 179
180 180 origsize = size
181 181
182 182 if size < 0:
183 183 size = remaining
184 184 else:
185 185 size = min(remaining, size)
186 186
187 187 result = object.__getattribute__(self, '_orig').readline(size)
188 188 remaining -= len(result)
189 189
190 190 self._writelog('readline(%d from %d) -> (%d) %s' % (
191 191 size, origsize, len(result), result))
192 192
193 193 object.__setattr__(self, '_closeafterrecvbytes', remaining)
194 194
195 195 if remaining <= 0:
196 196 self._writelog('read limit reached; closing socket')
197 197 self._sock.close()
198 198 # This is the easiest way to abort the current request.
199 199 raise Exception('connection closed after receiving N bytes')
200 200
201 201 return result
202 202
203 203 def write(self, data):
204 204 remaining = object.__getattribute__(self, '_closeaftersendbytes')
205 205
206 206 # No byte limit on this operation. Call original function.
207 207 if not remaining:
208 208 self._writelog('write(%d) -> %s' % (len(data), data))
209 209 result = object.__getattribute__(self, '_orig').write(data)
210 210 return result
211 211
212 212 if len(data) > remaining:
213 213 newdata = data[0:remaining]
214 214 else:
215 215 newdata = data
216 216
217 217 remaining -= len(newdata)
218 218
219 219 self._writelog('write(%d from %d) -> (%d) %s' % (
220 220 len(newdata), len(data), remaining, newdata))
221 221
222 222 result = object.__getattribute__(self, '_orig').write(newdata)
223 223
224 224 object.__setattr__(self, '_closeaftersendbytes', remaining)
225 225
226 226 if remaining <= 0:
227 227 self._writelog('write limit reached; closing socket')
228 228 self._sock.close()
229 229 raise Exception('connection closed after sending N bytes')
230 230
231 231 return result
232 232
233 233 def extsetup(ui):
234 234 # Change the base HTTP server class so various events can be performed.
235 235 # See SocketServer.BaseServer for how the specially named methods work.
236 236 class badserver(server.MercurialHTTPServer):
237 237 def __init__(self, ui, *args, **kwargs):
238 238 self._ui = ui
239 239 super(badserver, self).__init__(ui, *args, **kwargs)
240 240
241 241 # Need to inherit object so super() works.
242 242 class badrequesthandler(self.RequestHandlerClass, object):
243 243 def send_header(self, name, value):
244 244 # Make headers deterministic to facilitate testing.
245 245 if name.lower() == 'date':
246 246 value = 'Fri, 14 Apr 2017 00:00:00 GMT'
247 247 elif name.lower() == 'server':
248 248 value = 'badhttpserver'
249 249
250 250 return super(badrequesthandler, self).send_header(name,
251 251 value)
252 252
253 253 self.RequestHandlerClass = badrequesthandler
254 254
255 255 # Called to accept() a pending socket.
256 256 def get_request(self):
257 257 if self._ui.configbool('badserver', 'closebeforeaccept'):
258 258 self.socket.close()
259 259
260 260 # Tells the server to stop processing more requests.
261 261 self.__shutdown_request = True
262 262
263 263 # Simulate failure to stop processing this request.
264 264 raise socket.error('close before accept')
265 265
266 266 if self._ui.configbool('badserver', 'closeafteraccept'):
267 267 request, client_address = super(badserver, self).get_request()
268 268 request.close()
269 269 raise socket.error('close after accept')
270 270
271 271 return super(badserver, self).get_request()
272 272
273 273 # Does heavy lifting of processing a request. Invokes
274 274 # self.finish_request() which calls self.RequestHandlerClass() which
275 275 # is a hgweb.server._httprequesthandler.
276 276 def process_request(self, socket, address):
277 277 # Wrap socket in a proxy if we need to count bytes.
278 278 closeafterrecvbytes = self._ui.configint('badserver',
279 279 'closeafterrecvbytes')
280 280 closeaftersendbytes = self._ui.configint('badserver',
281 281 'closeaftersendbytes')
282 282
283 283 if closeafterrecvbytes or closeaftersendbytes:
284 284 socket = socketproxy(socket, self.errorlog,
285 285 closeafterrecvbytes=closeafterrecvbytes,
286 286 closeaftersendbytes=closeaftersendbytes)
287 287
288 288 return super(badserver, self).process_request(socket, address)
289 289
290 290 server.MercurialHTTPServer = badserver
General Comments 0
You need to be logged in to leave comments. Login now