##// END OF EJS Templates
tests: constant-fold a `pycompat.ispy3` in testlib/badserverext.py
Manuel Jacob -
r50193:425ca342 default
parent child Browse files
Show More
@@ -1,461 +1,456
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 close-before-accept
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 close-after-accept
22 22 If true, the server will close() the client socket immediately after
23 23 accept().
24 24
25 25 close-after-recv-bytes
26 26 If defined, close the client socket after receiving this many bytes.
27 27 (The value is a list, multiple values can use used to close a series of requests
28 28 request)
29 29
30 30 close-after-recv-patterns
31 31 If defined, the `close-after-recv-bytes` values only start counting after the
32 32 `read` operation that encountered the defined patterns.
33 33 (The value is a list, multiple values can use used to close a series of requests
34 34 request)
35 35
36 36 close-after-send-bytes
37 37 If defined, close the client socket after sending this many bytes.
38 38 (The value is a list, multiple values can use used to close a series of requests
39 39 request)
40 40
41 41 close-after-send-patterns
42 42 If defined, close the client socket after the configured regexp is seen.
43 43 (The value is a list, multiple values can use used to close a series of requests
44 44 request)
45 45 """
46 46
47 47
48 48 import re
49 49 import socket
50 50
51 51 from mercurial import (
52 pycompat,
53 52 registrar,
54 53 )
55 54
56 55 from mercurial.hgweb import server
57 56
58 57 configtable = {}
59 58 configitem = registrar.configitem(configtable)
60 59
61 60 configitem(
62 61 b'badserver',
63 62 b'close-after-accept',
64 63 default=False,
65 64 )
66 65 configitem(
67 66 b'badserver',
68 67 b'close-after-recv-bytes',
69 68 default=b'0',
70 69 )
71 70 configitem(
72 71 b'badserver',
73 72 b'close-after-recv-patterns',
74 73 default=b'',
75 74 )
76 75 configitem(
77 76 b'badserver',
78 77 b'close-after-send-bytes',
79 78 default=b'0',
80 79 )
81 80 configitem(
82 81 b'badserver',
83 82 b'close-after-send-patterns',
84 83 default=b'',
85 84 )
86 85 configitem(
87 86 b'badserver',
88 87 b'close-before-accept',
89 88 default=False,
90 89 )
91 90
92 91
93 92 class ConditionTracker:
94 93 def __init__(
95 94 self,
96 95 close_after_recv_bytes,
97 96 close_after_recv_patterns,
98 97 close_after_send_bytes,
99 98 close_after_send_patterns,
100 99 ):
101 100 self._all_close_after_recv_bytes = close_after_recv_bytes
102 101 self._all_close_after_recv_patterns = close_after_recv_patterns
103 102 self._all_close_after_send_bytes = close_after_send_bytes
104 103 self._all_close_after_send_patterns = close_after_send_patterns
105 104
106 105 self.target_recv_bytes = None
107 106 self.remaining_recv_bytes = None
108 107 self.recv_patterns = None
109 108 self.recv_data = b''
110 109 self.target_send_bytes = None
111 110 self.remaining_send_bytes = None
112 111 self.send_pattern = None
113 112 self.send_data = b''
114 113
115 114 def start_next_request(self):
116 115 """move to the next set of close condition"""
117 116 if self._all_close_after_recv_bytes:
118 117 self.target_recv_bytes = self._all_close_after_recv_bytes.pop(0)
119 118 self.remaining_recv_bytes = self.target_recv_bytes
120 119 else:
121 120 self.target_recv_bytes = None
122 121 self.remaining_recv_bytes = None
123 122
124 123 self.recv_data = b''
125 124 if self._all_close_after_recv_patterns:
126 125 self.recv_pattern = self._all_close_after_recv_patterns.pop(0)
127 126 else:
128 127 self.recv_pattern = None
129 128
130 129 if self._all_close_after_send_bytes:
131 130 self.target_send_bytes = self._all_close_after_send_bytes.pop(0)
132 131 self.remaining_send_bytes = self.target_send_bytes
133 132 else:
134 133 self.target_send_bytes = None
135 134 self.remaining_send_bytes = None
136 135
137 136 self.send_data = b''
138 137 if self._all_close_after_send_patterns:
139 138 self.send_pattern = self._all_close_after_send_patterns.pop(0)
140 139 else:
141 140 self.send_pattern = None
142 141
143 142 def might_close(self):
144 143 """True, if any processing will be needed"""
145 144 if self.remaining_recv_bytes is not None:
146 145 return True
147 146 if self.recv_pattern is not None:
148 147 return True
149 148 if self.remaining_send_bytes is not None:
150 149 return True
151 150 if self.send_pattern is not None:
152 151 return True
153 152 return False
154 153
155 154 def forward_write(self, obj, method, data, *args, **kwargs):
156 155 """call an underlying write function until condition are met
157 156
158 157 When the condition are met the socket is closed
159 158 """
160 159 remaining = self.remaining_send_bytes
161 160 pattern = self.send_pattern
162 161
163 162 orig = object.__getattribute__(obj, '_orig')
164 163 bmethod = method.encode('ascii')
165 164 func = getattr(orig, method)
166 165
167 166 if pattern:
168 167 self.send_data += data
169 168 pieces = pattern.split(self.send_data, maxsplit=1)
170 169 if len(pieces) > 1:
171 170 dropped = len(pieces[-1])
172 171 remaining = len(data) - dropped
173 172
174 173 if remaining:
175 174 remaining = max(0, remaining)
176 175
177 176 if not remaining:
178 177 newdata = data
179 178 else:
180 179 if remaining < len(data):
181 180 newdata = data[0:remaining]
182 181 else:
183 182 newdata = data
184 183 remaining -= len(newdata)
185 184 self.remaining_send_bytes = remaining
186 185
187 186 result = func(newdata, *args, **kwargs)
188 187
189 188 if remaining is None:
190 189 obj._writelog(b'%s(%d) -> %s' % (bmethod, len(data), data))
191 190 else:
192 191 msg = b'%s(%d from %d) -> (%d) %s'
193 192 msg %= (bmethod, len(newdata), len(data), remaining, newdata)
194 193 obj._writelog(msg)
195 194
196 195 if remaining is not None and remaining <= 0:
197 196 obj._writelog(b'write limit reached; closing socket')
198 197 object.__getattribute__(obj, '_cond_close')()
199 198 raise Exception('connection closed after sending N bytes')
200 199
201 200 return result
202 201
203 202 def forward_read(self, obj, method, size=-1):
204 203 """call an underlying read function until condition are met
205 204
206 205 When the condition are met the socket is closed
207 206 """
208 207 remaining = self.remaining_recv_bytes
209 208 pattern = self.recv_pattern
210 209
211 210 orig = object.__getattribute__(obj, '_orig')
212 211 bmethod = method.encode('ascii')
213 212 func = getattr(orig, method)
214 213
215 214 requested_size = size
216 215 actual_size = size
217 216
218 217 if pattern is None and remaining:
219 218 if size < 0:
220 219 actual_size = remaining
221 220 else:
222 221 actual_size = min(remaining, requested_size)
223 222
224 223 result = func(actual_size)
225 224
226 225 if pattern is None and remaining:
227 226 remaining -= len(result)
228 227 self.remaining_recv_bytes = remaining
229 228
230 229 if requested_size == 65537:
231 230 requested_repr = b'~'
232 231 else:
233 232 requested_repr = b'%d' % requested_size
234 233 if requested_size == actual_size:
235 234 msg = b'%s(%s) -> (%d) %s'
236 235 msg %= (bmethod, requested_repr, len(result), result)
237 236 else:
238 237 msg = b'%s(%d from %s) -> (%d) %s'
239 238 msg %= (bmethod, actual_size, requested_repr, len(result), result)
240 239 obj._writelog(msg)
241 240
242 241 if pattern is not None:
243 242 self.recv_data += result
244 243 if pattern.search(self.recv_data):
245 244 # start counting bytes starting with the next read
246 245 self.recv_pattern = None
247 246
248 247 if remaining is not None and remaining <= 0:
249 248 obj._writelog(b'read limit reached; closing socket')
250 249 obj._cond_close()
251 250
252 251 # This is the easiest way to abort the current request.
253 252 raise Exception('connection closed after receiving N bytes')
254 253
255 254 return result
256 255
257 256
258 257 # We can't adjust __class__ on a socket instance. So we define a proxy type.
259 258 class socketproxy:
260 259 __slots__ = ('_orig', '_logfp', '_cond')
261 260
262 261 def __init__(self, obj, logfp, condition_tracked):
263 262 object.__setattr__(self, '_orig', obj)
264 263 object.__setattr__(self, '_logfp', logfp)
265 264 object.__setattr__(self, '_cond', condition_tracked)
266 265
267 266 def __getattribute__(self, name):
268 267 if name in ('makefile', 'sendall', '_writelog', '_cond_close'):
269 268 return object.__getattribute__(self, name)
270 269
271 270 return getattr(object.__getattribute__(self, '_orig'), name)
272 271
273 272 def __delattr__(self, name):
274 273 delattr(object.__getattribute__(self, '_orig'), name)
275 274
276 275 def __setattr__(self, name, value):
277 276 setattr(object.__getattribute__(self, '_orig'), name, value)
278 277
279 278 def _writelog(self, msg):
280 279 msg = msg.replace(b'\r', b'\\r').replace(b'\n', b'\\n')
281 280
282 281 object.__getattribute__(self, '_logfp').write(msg)
283 282 object.__getattribute__(self, '_logfp').write(b'\n')
284 283 object.__getattribute__(self, '_logfp').flush()
285 284
286 285 def makefile(self, mode, bufsize):
287 286 f = object.__getattribute__(self, '_orig').makefile(mode, bufsize)
288 287
289 288 logfp = object.__getattribute__(self, '_logfp')
290 289 cond = object.__getattribute__(self, '_cond')
291 290
292 291 return fileobjectproxy(f, logfp, cond)
293 292
294 293 def sendall(self, data, flags=0):
295 294 cond = object.__getattribute__(self, '_cond')
296 295 return cond.forward_write(self, 'sendall', data, flags)
297 296
298 297 def _cond_close(self):
299 298 object.__getattribute__(self, '_orig').shutdown(socket.SHUT_RDWR)
300 299
301 300
302 301 # We can't adjust __class__ on socket._fileobject, so define a proxy.
303 302 class fileobjectproxy:
304 303 __slots__ = ('_orig', '_logfp', '_cond')
305 304
306 305 def __init__(self, obj, logfp, condition_tracked):
307 306 object.__setattr__(self, '_orig', obj)
308 307 object.__setattr__(self, '_logfp', logfp)
309 308 object.__setattr__(self, '_cond', condition_tracked)
310 309
311 310 def __getattribute__(self, name):
312 311 if name in (
313 312 '_close',
314 313 'read',
315 314 'readline',
316 315 'write',
317 316 '_writelog',
318 317 '_cond_close',
319 318 ):
320 319 return object.__getattribute__(self, name)
321 320
322 321 return getattr(object.__getattribute__(self, '_orig'), name)
323 322
324 323 def __delattr__(self, name):
325 324 delattr(object.__getattribute__(self, '_orig'), name)
326 325
327 326 def __setattr__(self, name, value):
328 327 setattr(object.__getattribute__(self, '_orig'), name, value)
329 328
330 329 def _writelog(self, msg):
331 330 msg = msg.replace(b'\r', b'\\r').replace(b'\n', b'\\n')
332 331
333 332 object.__getattribute__(self, '_logfp').write(msg)
334 333 object.__getattribute__(self, '_logfp').write(b'\n')
335 334 object.__getattribute__(self, '_logfp').flush()
336 335
337 336 def _close(self):
338 # Python 3 uses an io.BufferedIO instance. Python 2 uses some file
339 # object wrapper.
340 if pycompat.ispy3:
341 orig = object.__getattribute__(self, '_orig')
337 # We wrap an io.BufferedIO instance.
338 orig = object.__getattribute__(self, '_orig')
342 339
343 if hasattr(orig, 'raw'):
344 orig.raw._sock.shutdown(socket.SHUT_RDWR)
345 else:
346 self.close()
340 if hasattr(orig, 'raw'):
341 orig.raw._sock.shutdown(socket.SHUT_RDWR)
347 342 else:
348 self._sock.shutdown(socket.SHUT_RDWR)
343 self.close()
349 344
350 345 def read(self, size=-1):
351 346 cond = object.__getattribute__(self, '_cond')
352 347 return cond.forward_read(self, 'read', size)
353 348
354 349 def readline(self, size=-1):
355 350 cond = object.__getattribute__(self, '_cond')
356 351 return cond.forward_read(self, 'readline', size)
357 352
358 353 def write(self, data):
359 354 cond = object.__getattribute__(self, '_cond')
360 355 return cond.forward_write(self, 'write', data)
361 356
362 357 def _cond_close(self):
363 358 self._close()
364 359
365 360
366 361 def process_bytes_config(value):
367 362 parts = value.split(b',')
368 363 integers = [int(v) for v in parts if v]
369 364 return [v if v else None for v in integers]
370 365
371 366
372 367 def process_pattern_config(value):
373 368 patterns = []
374 369 for p in value.split(b','):
375 370 if not p:
376 371 p = None
377 372 else:
378 373 p = re.compile(p, re.DOTALL | re.MULTILINE)
379 374 patterns.append(p)
380 375 return patterns
381 376
382 377
383 378 def extsetup(ui):
384 379 # Change the base HTTP server class so various events can be performed.
385 380 # See SocketServer.BaseServer for how the specially named methods work.
386 381 class badserver(server.MercurialHTTPServer):
387 382 def __init__(self, ui, *args, **kwargs):
388 383 self._ui = ui
389 384 super(badserver, self).__init__(ui, *args, **kwargs)
390 385
391 386 all_recv_bytes = self._ui.config(
392 387 b'badserver', b'close-after-recv-bytes'
393 388 )
394 389 all_recv_bytes = process_bytes_config(all_recv_bytes)
395 390 all_recv_pattern = self._ui.config(
396 391 b'badserver', b'close-after-recv-patterns'
397 392 )
398 393 all_recv_pattern = process_pattern_config(all_recv_pattern)
399 394 all_send_bytes = self._ui.config(
400 395 b'badserver', b'close-after-send-bytes'
401 396 )
402 397 all_send_bytes = process_bytes_config(all_send_bytes)
403 398 all_send_patterns = self._ui.config(
404 399 b'badserver', b'close-after-send-patterns'
405 400 )
406 401 all_send_patterns = process_pattern_config(all_send_patterns)
407 402 self._cond = ConditionTracker(
408 403 all_recv_bytes,
409 404 all_recv_pattern,
410 405 all_send_bytes,
411 406 all_send_patterns,
412 407 )
413 408
414 409 # Need to inherit object so super() works.
415 410 class badrequesthandler(self.RequestHandlerClass, object):
416 411 def send_header(self, name, value):
417 412 # Make headers deterministic to facilitate testing.
418 413 if name.lower() == 'date':
419 414 value = 'Fri, 14 Apr 2017 00:00:00 GMT'
420 415 elif name.lower() == 'server':
421 416 value = 'badhttpserver'
422 417
423 418 return super(badrequesthandler, self).send_header(
424 419 name, value
425 420 )
426 421
427 422 self.RequestHandlerClass = badrequesthandler
428 423
429 424 # Called to accept() a pending socket.
430 425 def get_request(self):
431 426 if self._ui.configbool(b'badserver', b'close-before-accept'):
432 427 self.socket.close()
433 428
434 429 # Tells the server to stop processing more requests.
435 430 self.__shutdown_request = True
436 431
437 432 # Simulate failure to stop processing this request.
438 433 raise socket.error('close before accept')
439 434
440 435 if self._ui.configbool(b'badserver', b'close-after-accept'):
441 436 request, client_address = super(badserver, self).get_request()
442 437 request.close()
443 438 raise socket.error('close after accept')
444 439
445 440 return super(badserver, self).get_request()
446 441
447 442 # Does heavy lifting of processing a request. Invokes
448 443 # self.finish_request() which calls self.RequestHandlerClass() which
449 444 # is a hgweb.server._httprequesthandler.
450 445 def process_request(self, socket, address):
451 446 # Wrap socket in a proxy if we need to count bytes.
452 447 self._cond.start_next_request()
453 448
454 449 if self._cond.might_close():
455 450 socket = socketproxy(
456 451 socket, self.errorlog, condition_tracked=self._cond
457 452 )
458 453
459 454 return super(badserver, self).process_request(socket, address)
460 455
461 456 server.MercurialHTTPServer = badserver
General Comments 0
You need to be logged in to leave comments. Login now