mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-23 08:11:00 +00:00
Unify pause and disconnect event frameworks.
This commit is contained in:
Aldo Cortesi
parent
7810ee475d
commit
5052a6d34e
@ -125,13 +125,6 @@ class ValueLiteral:
|
||||
return self.val
|
||||
|
||||
|
||||
class ValueNakedLiteral(ValueLiteral):
|
||||
@classmethod
|
||||
def expr(klass):
|
||||
e = v_naked_literal.copy()
|
||||
return e.setParseAction(lambda x: klass(*x))
|
||||
|
||||
|
||||
class ValueGenerate:
|
||||
UNITS = dict(
|
||||
b = 1024**0,
|
||||
@ -232,19 +225,19 @@ class _Pause:
|
||||
class PauseBefore(_Pause):
|
||||
sub = "b"
|
||||
def mod_response(self, settings, r):
|
||||
r.pauses.append((0, self.value))
|
||||
r.actions.append((0, "pause", self.value))
|
||||
|
||||
|
||||
class PauseAfter(_Pause):
|
||||
sub = "a"
|
||||
def mod_response(self, settings, r):
|
||||
r.pauses.append((sys.maxint, self.value))
|
||||
r.actions.append((sys.maxint, "pause", self.value))
|
||||
|
||||
|
||||
class PauseRandom(_Pause):
|
||||
sub = "r"
|
||||
def mod_response(self, settings, r):
|
||||
r.pauses.append(("random", self.value))
|
||||
r.actions.append(("random", "pause", self.value))
|
||||
|
||||
|
||||
|
||||
@ -261,13 +254,13 @@ class _Disconnect:
|
||||
class DisconnectBefore(_Disconnect):
|
||||
sub = "b"
|
||||
def mod_response(self, settings, r):
|
||||
r.pauses.append((0, self.value))
|
||||
r.actions.append((0, "disconnect"))
|
||||
|
||||
|
||||
class DisconnectRandom(_Disconnect):
|
||||
sub = "r"
|
||||
def mod_response(self, settings, r):
|
||||
r.pauses.append(("random", self.value))
|
||||
r.actions.append(("random", "disconnect"))
|
||||
|
||||
|
||||
class Header:
|
||||
@ -312,7 +305,7 @@ class Code:
|
||||
|
||||
BLOCKSIZE = 1024
|
||||
class Response:
|
||||
comps = [
|
||||
comps = (
|
||||
Body,
|
||||
Header,
|
||||
PauseBefore,
|
||||
@ -320,14 +313,14 @@ class Response:
|
||||
PauseRandom,
|
||||
DisconnectBefore,
|
||||
DisconnectRandom,
|
||||
]
|
||||
)
|
||||
version = "HTTP/1.1"
|
||||
code = 200
|
||||
msg = LiteralGenerator(http.RESPONSES[code])
|
||||
body = LiteralGenerator("OK")
|
||||
def __init__(self):
|
||||
self.headers = []
|
||||
self.pauses = []
|
||||
self.actions = []
|
||||
|
||||
def get_header(self, hdr):
|
||||
for k, v in self.headers:
|
||||
@ -374,30 +367,33 @@ class Response:
|
||||
else:
|
||||
tornado.ioloop.IOLoop.instance().add_timeout(time.time() + s, callback)
|
||||
|
||||
def write_values(self, fp, vals, pauses, disconnect, sofar=0, skip=0, blocksize=BLOCKSIZE):
|
||||
if disconnect == "before":
|
||||
fp.finish()
|
||||
return
|
||||
def write_values(self, fp, vals, actions, sofar=0, skip=0, blocksize=BLOCKSIZE):
|
||||
while vals:
|
||||
part = vals.pop()
|
||||
for i in range(skip, len(part), blocksize):
|
||||
d = part[i:i+blocksize]
|
||||
if pauses and pauses[-1][0] < (sofar + len(d)):
|
||||
p = pauses.pop()
|
||||
if actions and actions[-1][0] < (sofar + len(d)):
|
||||
p = actions.pop()
|
||||
offset = p[0]-sofar
|
||||
vals.append(part)
|
||||
def pause_callback():
|
||||
self.write_values(
|
||||
fp, vals, pauses, disconnect,
|
||||
sofar=sofar+offset,
|
||||
skip=i+offset,
|
||||
blocksize=blocksize
|
||||
)
|
||||
def flushed_callback():
|
||||
# Data has been flushed, set the timeout.
|
||||
self.add_timeout(p[1], pause_callback)
|
||||
fp.write(d[:offset], callback=flushed_callback)
|
||||
return
|
||||
if p[1] == "pause":
|
||||
def pause_callback():
|
||||
self.write_values(
|
||||
fp, vals, actions,
|
||||
sofar=sofar+offset,
|
||||
skip=i+offset,
|
||||
blocksize=blocksize
|
||||
)
|
||||
def flushed_callback():
|
||||
# Data has been flushed, set the timeout.
|
||||
self.add_timeout(p[2], pause_callback)
|
||||
fp.write(d[:offset], callback=flushed_callback)
|
||||
return
|
||||
elif p[1] == "disconnect":
|
||||
fp.write(d[:offset])
|
||||
fp.finish()
|
||||
fp.connection.stream.close()
|
||||
return
|
||||
fp.write(d)
|
||||
sofar += len(d)
|
||||
skip = 0
|
||||
@ -431,9 +427,9 @@ class Response:
|
||||
self.body
|
||||
])
|
||||
vals.reverse()
|
||||
pauses = self.ready_randoms(self.length(), self.pauses)
|
||||
pauses.reverse()
|
||||
return self.write_values(fp, vals, pauses, None)
|
||||
actions = self.ready_randoms(self.length(), self.actions)
|
||||
actions.reverse()
|
||||
return self.write_values(fp, vals, actions)
|
||||
|
||||
def __str__(self):
|
||||
parts = [
|
||||
|
||||
@ -4,7 +4,17 @@ from libpathod import rparse
|
||||
|
||||
rparse.TESTING = True
|
||||
|
||||
|
||||
class Sponge:
|
||||
def __getattr__(self, x):
|
||||
return Sponge()
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
|
||||
class DummyRequest(StringIO.StringIO):
|
||||
connection = Sponge()
|
||||
def write(self, d, callback=None):
|
||||
StringIO.StringIO.write(self, d)
|
||||
if callback:
|
||||
@ -54,6 +64,7 @@ class uMisc(libpry.AutoTree):
|
||||
|
||||
def test_file_value(self):
|
||||
v = rparse.Value.parseString("<'one two'")[0]
|
||||
assert str(v)
|
||||
assert v.path == "one two"
|
||||
|
||||
v = rparse.Value.parseString("<path")[0]
|
||||
@ -147,6 +158,9 @@ class uMisc(libpry.AutoTree):
|
||||
|
||||
|
||||
class uDisconnects(libpry.AutoTree):
|
||||
def test_parse(self):
|
||||
assert rparse.parse({}, "400:db")
|
||||
|
||||
def test_before(self):
|
||||
e = rparse.DisconnectBefore.expr()
|
||||
v = e.parseString("db")[0]
|
||||
@ -185,12 +199,14 @@ class uPauses(libpry.AutoTree):
|
||||
|
||||
|
||||
class uparse(libpry.AutoTree):
|
||||
|
||||
def test_parse_err(self):
|
||||
libpry.raises(rparse.ParseException, rparse.parse, {}, "400:msg,b:")
|
||||
try:
|
||||
rparse.parse({}, "400'msg':b:")
|
||||
except rparse.ParseException, v:
|
||||
assert v.marked()
|
||||
assert str(v)
|
||||
|
||||
def test_parse_header(self):
|
||||
r = rparse.parse({}, '400:h"foo"="bar"')
|
||||
@ -198,15 +214,15 @@ class uparse(libpry.AutoTree):
|
||||
|
||||
def test_parse_pause_before(self):
|
||||
r = rparse.parse({}, "400:pb10")
|
||||
assert (0, 10) in r.pauses
|
||||
assert (0, "pause", 10) in r.actions
|
||||
|
||||
def test_parse_pause_after(self):
|
||||
r = rparse.parse({}, "400:pa10")
|
||||
assert (sys.maxint, 10) in r.pauses
|
||||
assert (sys.maxint, "pause", 10) in r.actions
|
||||
|
||||
def test_parse_pause_random(self):
|
||||
r = rparse.parse({}, "400:pr10")
|
||||
assert ("random", 10) in r.pauses
|
||||
assert ("random", "pause", 10) in r.actions
|
||||
|
||||
|
||||
class uResponse(libpry.AutoTree):
|
||||
@ -240,7 +256,7 @@ class uResponse(libpry.AutoTree):
|
||||
r = self.dummy_response()
|
||||
s = DummyRequest()
|
||||
tst = "foo"*100
|
||||
r.write_values(s, [tst], [], "before", blocksize=5)
|
||||
r.write_values(s, [tst], [(0, "disconnect")], blocksize=5)
|
||||
assert not s.getvalue()
|
||||
|
||||
def test_write_values(self):
|
||||
@ -248,7 +264,7 @@ class uResponse(libpry.AutoTree):
|
||||
r = rparse.parse({}, "400'msg'")
|
||||
|
||||
s = DummyRequest()
|
||||
r.write_values(s, [tst], [], None)
|
||||
r.write_values(s, [tst], [])
|
||||
assert s.getvalue() == tst
|
||||
|
||||
def test_write_values_pauses(self):
|
||||
@ -257,18 +273,18 @@ class uResponse(libpry.AutoTree):
|
||||
|
||||
for i in range(2, 10):
|
||||
s = DummyRequest()
|
||||
r.write_values(s, [tst], [(2, 0), (1, 0)], None, blocksize=i)
|
||||
r.write_values(s, [tst], [(2, "pause", 0), (1, "pause", 0)], blocksize=i)
|
||||
assert s.getvalue() == tst
|
||||
|
||||
for i in range(2, 10):
|
||||
s = DummyRequest()
|
||||
r.write_values(s, [tst], [(1, 0)], None, blocksize=i)
|
||||
r.write_values(s, [tst], [(1, "pause", 0)], blocksize=i)
|
||||
assert s.getvalue() == tst
|
||||
|
||||
tst = ["".join(str(i) for i in range(10))]*5
|
||||
for i in range(2, 10):
|
||||
s = DummyRequest()
|
||||
r.write_values(s, tst[:], [(1, 0)], None, blocksize=i)
|
||||
r.write_values(s, tst[:], [(1, "pause", 0)], blocksize=i)
|
||||
assert s.getvalue() == "".join(tst)
|
||||
|
||||
def test_render(self):
|
||||
|
||||
Loading…
Reference in New Issue
Block a user