From 4f7eef6029f038425c4e10da70ce5f65e25d77c7 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Sat, 8 Dec 2018 13:07:12 +0200
Subject: [PATCH] Fix uploading and downloading avatars

---
 maubot/management/api/__init__.py     | 2 +-
 maubot/management/api/client_proxy.py | 9 ++-------
 maubot/management/api/middleware.py   | 4 +---
 maubot/server.py                      | 2 +-
 4 files changed, 5 insertions(+), 12 deletions(-)

diff --git a/maubot/management/api/__init__.py b/maubot/management/api/__init__.py
index 93e994d..ceb86ef 100644
--- a/maubot/management/api/__init__.py
+++ b/maubot/management/api/__init__.py
@@ -32,7 +32,7 @@ from .log import stop_all as stop_log_sockets, init as init_log_listener
 def init(cfg: Config, loop: AbstractEventLoop) -> web.Application:
     set_config(cfg)
     set_loop(loop)
-    app = web.Application(loop=loop, middlewares=[auth, error])
+    app = web.Application(loop=loop, middlewares=[auth, error], client_max_size=100*1024*1024)
     app.add_routes(routes)
     return app
 
diff --git a/maubot/management/api/client_proxy.py b/maubot/management/api/client_proxy.py
index b6f5787..8dbb650 100644
--- a/maubot/management/api/client_proxy.py
+++ b/maubot/management/api/client_proxy.py
@@ -14,7 +14,6 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 from aiohttp import web, client as http
-
 from ...client import Client
 from .base import routes
 from .responses import resp
@@ -44,15 +43,11 @@ async def proxy(request: web.Request) -> web.StreamResponse:
             headers["X-Forwarded-For"] = f"{host}:{port}"
 
     data = await request.read()
-    chunked = PROXY_CHUNK_SIZE if not data else None
     async with http.request(request.method, f"{client.homeserver}/{path}", headers=headers,
-                            params=query, chunked=chunked, data=data) as proxy_resp:
+                            params=query, data=data) as proxy_resp:
         response = web.StreamResponse(status=proxy_resp.status, headers=proxy_resp.headers)
         await response.prepare(request)
-        content = proxy_resp.content
-        chunk = await content.read(PROXY_CHUNK_SIZE)
-        while chunk:
+        async for chunk in proxy_resp.content.iter_chunked(PROXY_CHUNK_SIZE):
             await response.write(chunk)
-            chunk = await content.read(PROXY_CHUNK_SIZE)
         await response.write_eof()
         return response
diff --git a/maubot/management/api/middleware.py b/maubot/management/api/middleware.py
index 538dba5..7406b61 100644
--- a/maubot/management/api/middleware.py
+++ b/maubot/management/api/middleware.py
@@ -23,6 +23,7 @@ from .auth import check_token
 from .base import get_config
 
 Handler = Callable[[web.Request], Awaitable[web.Response]]
+log = logging.getLogger("maubot.server")
 
 
 @web.middleware
@@ -36,9 +37,6 @@ async def auth(request: web.Request, handler: Handler) -> web.Response:
     return await handler(request)
 
 
-log = logging.getLogger("maubot.server")
-
-
 @web.middleware
 async def error(request: web.Request, handler: Handler) -> web.Response:
     try:
diff --git a/maubot/server.py b/maubot/server.py
index deb17e6..35682ee 100644
--- a/maubot/server.py
+++ b/maubot/server.py
@@ -38,7 +38,7 @@ class MaubotServer:
 
     def __init__(self, config: Config, loop: asyncio.AbstractEventLoop) -> None:
         self.loop = loop or asyncio.get_event_loop()
-        self.app = web.Application(loop=self.loop)
+        self.app = web.Application(loop=self.loop, client_max_size=100*1024*1024)
         self.config = config
 
         as_path = PathBuilder(config["server.appservice_base_path"])