25 close-after-recv-bytes |
25 close-after-recv-bytes |
26 If defined, close the client socket after receiving this many bytes. |
26 If defined, close the client socket after receiving this many bytes. |
27 (The value is a list, multiple values can use used to close a series of requests |
27 (The value is a list, multiple values can use used to close a series of requests |
28 request) |
28 request) |
29 |
29 |
|
30 close-after-recv-patterns |
|
31 If defined, the `close-after-recv-bytes` values only start counting after the |
|
32 `read` operation that encountered the defined patterns. |
|
33 (The value is a list, multiple values can use used to close a series of requests |
|
34 request) |
|
35 |
30 close-after-send-bytes |
36 close-after-send-bytes |
31 If defined, close the client socket after sending this many bytes. |
37 If defined, close the client socket after sending this many bytes. |
32 (The value is a list, multiple values can use used to close a series of requests |
38 (The value is a list, multiple values can use used to close a series of requests |
33 request) |
39 request) |
34 |
40 |
82 |
93 |
83 class ConditionTracker(object): |
94 class ConditionTracker(object): |
84 def __init__( |
95 def __init__( |
85 self, |
96 self, |
86 close_after_recv_bytes, |
97 close_after_recv_bytes, |
|
98 close_after_recv_patterns, |
87 close_after_send_bytes, |
99 close_after_send_bytes, |
88 close_after_send_patterns, |
100 close_after_send_patterns, |
89 ): |
101 ): |
90 self._all_close_after_recv_bytes = close_after_recv_bytes |
102 self._all_close_after_recv_bytes = close_after_recv_bytes |
|
103 self._all_close_after_recv_patterns = close_after_recv_patterns |
91 self._all_close_after_send_bytes = close_after_send_bytes |
104 self._all_close_after_send_bytes = close_after_send_bytes |
92 self._all_close_after_send_patterns = close_after_send_patterns |
105 self._all_close_after_send_patterns = close_after_send_patterns |
93 |
106 |
94 self.target_recv_bytes = None |
107 self.target_recv_bytes = None |
95 self.remaining_recv_bytes = None |
108 self.remaining_recv_bytes = None |
|
109 self.recv_patterns = None |
|
110 self.recv_data = b'' |
96 self.target_send_bytes = None |
111 self.target_send_bytes = None |
97 self.remaining_send_bytes = None |
112 self.remaining_send_bytes = None |
98 self.send_pattern = None |
113 self.send_pattern = None |
99 self.send_data = b'' |
114 self.send_data = b'' |
100 |
115 |
105 self.remaining_recv_bytes = self.target_recv_bytes |
120 self.remaining_recv_bytes = self.target_recv_bytes |
106 else: |
121 else: |
107 self.target_recv_bytes = None |
122 self.target_recv_bytes = None |
108 self.remaining_recv_bytes = None |
123 self.remaining_recv_bytes = None |
109 |
124 |
|
125 self.recv_data = b'' |
|
126 if self._all_close_after_recv_patterns: |
|
127 self.recv_pattern = self._all_close_after_recv_patterns.pop(0) |
|
128 else: |
|
129 self.recv_pattern = None |
|
130 |
110 if self._all_close_after_send_bytes: |
131 if self._all_close_after_send_bytes: |
111 self.target_send_bytes = self._all_close_after_send_bytes.pop(0) |
132 self.target_send_bytes = self._all_close_after_send_bytes.pop(0) |
112 self.remaining_send_bytes = self.target_send_bytes |
133 self.remaining_send_bytes = self.target_send_bytes |
113 else: |
134 else: |
114 self.target_send_bytes = None |
135 self.target_send_bytes = None |
121 self.send_pattern = None |
142 self.send_pattern = None |
122 |
143 |
123 def might_close(self): |
144 def might_close(self): |
124 """True, if any processing will be needed""" |
145 """True, if any processing will be needed""" |
125 if self.remaining_recv_bytes is not None: |
146 if self.remaining_recv_bytes is not None: |
|
147 return True |
|
148 if self.recv_pattern is not None: |
126 return True |
149 return True |
127 if self.remaining_send_bytes is not None: |
150 if self.remaining_send_bytes is not None: |
128 return True |
151 return True |
129 if self.send_pattern is not None: |
152 if self.send_pattern is not None: |
130 return True |
153 return True |
182 """call an underlying read function until condition are met |
205 """call an underlying read function until condition are met |
183 |
206 |
184 When the condition are met the socket is closed |
207 When the condition are met the socket is closed |
185 """ |
208 """ |
186 remaining = self.remaining_recv_bytes |
209 remaining = self.remaining_recv_bytes |
|
210 pattern = self.recv_pattern |
187 |
211 |
188 orig = object.__getattribute__(obj, '_orig') |
212 orig = object.__getattribute__(obj, '_orig') |
189 bmethod = method.encode('ascii') |
213 bmethod = method.encode('ascii') |
190 func = getattr(orig, method) |
214 func = getattr(orig, method) |
191 |
215 |
192 requested_size = size |
216 requested_size = size |
193 actual_size = size |
217 actual_size = size |
194 |
218 |
195 if remaining: |
219 if pattern is None and remaining: |
196 if size < 0: |
220 if size < 0: |
197 actual_size = remaining |
221 actual_size = remaining |
198 else: |
222 else: |
199 actual_size = min(remaining, requested_size) |
223 actual_size = min(remaining, requested_size) |
200 |
224 |
201 result = func(actual_size) |
225 result = func(actual_size) |
202 |
226 |
203 if remaining: |
227 if pattern is None and remaining: |
204 remaining -= len(result) |
228 remaining -= len(result) |
205 self.remaining_recv_bytes = remaining |
229 self.remaining_recv_bytes = remaining |
206 |
230 |
207 if requested_size == 65537: |
231 if requested_size == 65537: |
208 requested_repr = b'~' |
232 requested_repr = b'~' |
213 msg %= (bmethod, requested_repr, len(result), result) |
237 msg %= (bmethod, requested_repr, len(result), result) |
214 else: |
238 else: |
215 msg = b'%s(%d from %s) -> (%d) %s' |
239 msg = b'%s(%d from %s) -> (%d) %s' |
216 msg %= (bmethod, actual_size, requested_repr, len(result), result) |
240 msg %= (bmethod, actual_size, requested_repr, len(result), result) |
217 obj._writelog(msg) |
241 obj._writelog(msg) |
|
242 |
|
243 if pattern is not None: |
|
244 self.recv_data += result |
|
245 if pattern.search(self.recv_data): |
|
246 # start counting bytes starting with the next read |
|
247 self.recv_pattern = None |
218 |
248 |
219 if remaining is not None and remaining <= 0: |
249 if remaining is not None and remaining <= 0: |
220 obj._writelog(b'read limit reached; closing socket') |
250 obj._writelog(b'read limit reached; closing socket') |
221 obj._cond_close() |
251 obj._cond_close() |
222 |
252 |
361 |
391 |
362 all_recv_bytes = self._ui.config( |
392 all_recv_bytes = self._ui.config( |
363 b'badserver', b'close-after-recv-bytes' |
393 b'badserver', b'close-after-recv-bytes' |
364 ) |
394 ) |
365 all_recv_bytes = process_bytes_config(all_recv_bytes) |
395 all_recv_bytes = process_bytes_config(all_recv_bytes) |
|
396 all_recv_pattern = self._ui.config( |
|
397 b'badserver', b'close-after-recv-patterns' |
|
398 ) |
|
399 all_recv_pattern = process_pattern_config(all_recv_pattern) |
366 all_send_bytes = self._ui.config( |
400 all_send_bytes = self._ui.config( |
367 b'badserver', b'close-after-send-bytes' |
401 b'badserver', b'close-after-send-bytes' |
368 ) |
402 ) |
369 all_send_bytes = process_bytes_config(all_send_bytes) |
403 all_send_bytes = process_bytes_config(all_send_bytes) |
370 all_send_patterns = self._ui.config( |
404 all_send_patterns = self._ui.config( |
371 b'badserver', b'close-after-send-patterns' |
405 b'badserver', b'close-after-send-patterns' |
372 ) |
406 ) |
373 all_send_patterns = process_pattern_config(all_send_patterns) |
407 all_send_patterns = process_pattern_config(all_send_patterns) |
374 self._cond = ConditionTracker( |
408 self._cond = ConditionTracker( |
375 all_recv_bytes, |
409 all_recv_bytes, |
|
410 all_recv_pattern, |
376 all_send_bytes, |
411 all_send_bytes, |
377 all_send_patterns, |
412 all_send_patterns, |
378 ) |
413 ) |
379 |
414 |
380 # Need to inherit object so super() works. |
415 # Need to inherit object so super() works. |