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 |
# |
|
|
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 |
|
|
|
344 |
|
|
|
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