Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-23 00:01:36 +00:00
Code Issues Projects Releases Wiki Activity

tcp layer: handle connection error

Browse Source
This commit is contained in:
Maximilian Hils 2017-06-28 17:25:14 +02:00
parent 4bca88608b
commit c4824f941d
2 changed files with 38 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mitmproxy/proxy/protocol2/tcp.py
View File

@ -33,6 +33,8 @@ class TCPLayer(Layer):
self.flow.error = flow.Error(err) self.flow.error = flow.Error(err)
yield commands.Hook("tcp_error", self.flow) yield commands.Hook("tcp_error", self.flow)
yield commands.CloseConnection(self.context.client) yield commands.CloseConnection(self.context.client)
self._handle_event = self.done
return
self._handle_event = self.relay_messages self._handle_event = self.relay_messages
_handle_event = start _handle_event = start

42
mitmproxy/proxy/protocol2/test/test_tcp.py
View File

@ -21,6 +21,20 @@ def test_open_connection(tctx):
) )
def test_open_connection_err(tctx):
f = tutils.Placeholder()
assert (
tutils.playbook(tcp.TCPLayer(tctx))
<< commands.Hook("tcp_start", f)
>> events.HookReply(-1, None)
<< commands.OpenConnection(tctx.server)
>> events.OpenConnectionReply(-1, "Connect call failed")
<< commands.Hook("tcp_error", f)
>> events.HookReply(-1, None)
<< commands.CloseConnection(tctx.client)
)
def test_simple(tctx): def test_simple(tctx):
"""open connection, receive data, send it to peer""" """open connection, receive data, send it to peer"""
f = tutils.Placeholder() f = tutils.Placeholder()
@ -34,18 +48,20 @@ def test_simple(tctx):
>> events.OpenConnectionReply(-1, None) >> events.OpenConnectionReply(-1, None)
>> events.DataReceived(tctx.client, b"hello!") >> events.DataReceived(tctx.client, b"hello!")
<< commands.Hook("tcp_message", f) << commands.Hook("tcp_message", f)
)
assert f().messages[0].content == b"hello!"
assert (
playbook
>> events.HookReply(-1, None) >> events.HookReply(-1, None)
<< commands.SendData(tctx.server, b"hello!") << commands.SendData(tctx.server, b"hello!")
>> events.DataReceived(tctx.server, b"hi")
<< commands.Hook("tcp_message", f)
>> events.HookReply(-1, None)
<< commands.SendData(tctx.client, b"hi")
>> events.ConnectionClosed(tctx.server) >> events.ConnectionClosed(tctx.server)
<< commands.CloseConnection(tctx.client) << commands.CloseConnection(tctx.client)
<< commands.Hook("tcp_end", f) << commands.Hook("tcp_end", f)
>> events.HookReply(-1, None) >> events.HookReply(-1, None)
>> events.ConnectionClosed(tctx.client)
<< None << None
) )
assert len(f().messages) == 2
def test_simple_explicit(tctx): def test_simple_explicit(tctx):
@ -65,8 +81,7 @@ def test_simple_explicit(tctx):
assert flow.messages[0].content == b"hello!" assert flow.messages[0].content == b"hello!"
send, = layer.handle_event(events.HookReply(tcp_msg, None)) send, = layer.handle_event(events.HookReply(tcp_msg, None))
assert tutils._eq(send, commands.SendData(tctx.server, b"hello!s" assert tutils._eq(send, commands.SendData(tctx.server, b"hello!"))
b""))
close, tcp_end = layer.handle_event(events.ConnectionClosed(tctx.server)) close, tcp_end = layer.handle_event(events.ConnectionClosed(tctx.server))
assert tutils._eq(close, commands.CloseConnection(tctx.client)) assert tutils._eq(close, commands.CloseConnection(tctx.client))
assert tutils._eq(tcp_end, commands.Hook("tcp_end", flow)) assert tutils._eq(tcp_end, commands.Hook("tcp_end", flow))
@ -120,3 +135,18 @@ def test_receive_data_before_server_connected(tctx):
>> events.HookReply(-1, None) >> events.HookReply(-1, None)
<< commands.SendData(tctx.server, b"hello!") << commands.SendData(tctx.server, b"hello!")
) )
def test_receive_data_after_server_disconnected(tctx):
"""
this should just be discarded.
"""
assert (
tutils.playbook(tcp.TCPLayer(tctx, True))
<< commands.OpenConnection(tctx.server)
>> events.OpenConnectionReply(-1, None)
>> events.ConnectionClosed(tctx.server)
<< commands.CloseConnection(tctx.client)
>> events.DataReceived(tctx.client, b"i'm late")
<< None
)