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

async server: implement close command

Browse Source
This commit is contained in:
Maximilian Hils 2017-06-27 20:05:14 +02:00
parent b14b487cbe
commit 35475f6719
1 changed files with 19 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
mitmproxy/proxy/protocol2/server/server_async.py
View File

@ -26,8 +26,8 @@ class ConnectionHandler:
self.client = Client(addr) self.client = Client(addr)
self.context = Context(self.client) self.context = Context(self.client)
# self.layer = ReverseProxy(self.context, ("example.com", 443)) # self.layer = ReverseProxy(self.context, ("localhost", 443))
self.layer = ReverseProxy(self.context, ("example.com", 80)) self.layer = ReverseProxy(self.context, ("localhost", 80))
self.transports: MutableMapping[Connection, StreamIO] = { self.transports: MutableMapping[Connection, StreamIO] = {
self.client: StreamIO(reader, writer) self.client: StreamIO(reader, writer)
@ -41,15 +41,19 @@ class ConnectionHandler:
print("client connection done, closing transports!") print("client connection done, closing transports!")
for connection in list(self.transports): if self.transports:
await self.close(connection) await asyncio.wait([
self.close_connection(x)
for x in self.transports
])
# TODO: teardown all other conns.
print("transports closed!") print("transports closed!")
async def close(self, connection): async def close_connection(self, connection):
print("Closing", connection) io = self.transports.pop(connection, None)
io = self.transports.pop(connection) if not io:
print(f"Already closed: {connection}")
print(f"Closing {connection}")
try: try:
await io.w.drain() await io.w.drain()
io.w.write_eof() io.w.write_eof()
@ -69,7 +73,7 @@ class ConnectionHandler:
else: else:
connection.connected = False connection.connected = False
if connection in self.transports: if connection in self.transports:
await self.close(connection) await self.close_connection(connection)
await self.server_event(events.ConnectionClosed(connection)) await self.server_event(events.ConnectionClosed(connection))
break break
@ -97,10 +101,14 @@ class ConnectionHandler:
# TODO: pass to master here. # TODO: pass to master here.
print(f"~ {command.name}: {command.data}") print(f"~ {command.name}: {command.data}")
asyncio.ensure_future( asyncio.ensure_future(
self.server_event(events.HookReply(command, "hook reply")) self.server_event(events.HookReply(command, None))
)
elif isinstance(command, commands.CloseConnection):
asyncio.ensure_future(
self.close_connection(command.connection)
) )
else: else:
raise NotImplementedError("Unexpected event: {}".format(command)) raise NotImplementedError(f"Unexpected event: {command}")
print("#>") print("#>")