From 82104266fec249073346662d6ce75afadc804ce1 Mon Sep 17 00:00:00 2001
From: miruka <miruka@disroot.org>
Date: Mon, 4 Nov 2019 15:18:01 -0400
Subject: [PATCH] Encrypt uploads in a thread

---
 TODO.md                     |  6 ++----
 src/python/matrix_client.py | 31 ++++++++++++++++++++++++-------
 src/python/media_cache.py   |  2 +-
 3 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/TODO.md b/TODO.md
index 35f5c2a6..595dc713 100644
--- a/TODO.md
+++ b/TODO.md
@@ -1,9 +1,7 @@
 - Media
-  - Encrypt file for upload in thread
   - Cache our own uploads
-  - Uploading progress (+local echo)
-    - Deduplicate uploads
-  - Loading progress bar
+  - Uploading progress bar (+local echo)
+  - Image loading progress bar
   - Downloading
     - Bottom/top bar
 
diff --git a/src/python/matrix_client.py b/src/python/matrix_client.py
index 29b67a8e..5ca84f08 100644
--- a/src/python/matrix_client.py
+++ b/src/python/matrix_client.py
@@ -1,4 +1,5 @@
 import asyncio
+import functools
 import html
 import inspect
 import io
@@ -12,20 +13,24 @@ from datetime import datetime
 from functools import partial
 from pathlib import Path
 from types import ModuleType
-from typing import Any, DefaultDict, Dict, Optional, Set, Tuple, Type, Union
+from typing import (
+    Any, BinaryIO, DefaultDict, Dict, Optional, Set, Tuple, Type, Union
+)
 from uuid import uuid4
 
-import nio
-from nio.crypto.attachments import encrypt_attachment
 from PIL import Image as PILImage
 from pymediainfo import MediaInfo
 
+import nio
+
 from . import __about__, utils
 from .html_filter import HTML_FILTER
 from .models.items import Account, Event, Member, Room, TypeSpecifier
 from .models.model_store import ModelStore
 from .pyotherside_events import AlertRequested
 
+CryptDict = Dict[str, Any]
+
 
 @dataclass
 class UploadError(Exception):
@@ -419,9 +424,19 @@ class MatrixClient(nio.AsyncClient):
         self.models.pop((Member, room_id), None)
 
 
+    async def encrypt_attachment(self, data: bytes) -> Tuple[bytes, CryptDict]:
+        func = functools.partial(
+            nio.crypto.attachments.encrypt_attachment,
+            data,
+        )
+
+        # Run in a separate thread
+        return await asyncio.get_event_loop().run_in_executor(None, func)
+
+
     async def upload_thumbnail(
         self, path: Union[Path, str], encrypt: bool = False,
-    ) -> Tuple[str, Dict[str, Any], Dict[str, Any]]:
+    ) -> Tuple[str, Dict[str, Any], CryptDict]:
 
         png_modes = ("1", "L", "P", "RGBA")
 
@@ -449,7 +464,7 @@ class MatrixClient(nio.AsyncClient):
                 data = out.getvalue()
 
                 if encrypt:
-                    data, crypt_dict = encrypt_attachment(data)
+                    data, crypt_dict = await self.encrypt_attachment(data)
                     upload_mime      = "application/octet-stream"
                 else:
                     crypt_dict, upload_mime = {}, mime
@@ -471,13 +486,15 @@ class MatrixClient(nio.AsyncClient):
 
 
     async def upload_file(self, path: Union[Path, str], encrypt: bool = False,
-                         ) -> Tuple[str, str, Dict[str, Any]]:
+                         ) -> Tuple[str, str, CryptDict]:
         with open(path, "rb") as file:
             mime = utils.guess_mime(file)
             file.seek(0, 0)
 
+            data: Union[BinaryIO, bytes]
+
             if encrypt:
-                data, crypt_dict = encrypt_attachment(file.read())
+                data, crypt_dict = await self.encrypt_attachment(file.read())
                 upload_mime      = "application/octet-stream"
             else:
                 data, crypt_dict, upload_mime = file, {}, mime
diff --git a/src/python/media_cache.py b/src/python/media_cache.py
index 0a17d048..5ecbf9e2 100644
--- a/src/python/media_cache.py
+++ b/src/python/media_cache.py
@@ -12,7 +12,7 @@ from PIL import Image as PILImage
 
 import nio
 
-from .matrix_client import MatrixClient
+from .matrix_client import MatrixClient, CryptDict
 
 CryptDict = Optional[Dict[str, Any]]
 Size      = Tuple[int, int]