Matrix-Mirrors/pantalaimon
1
0
Fork 0
mirror of https://github.com/matrix-org/pantalaimon.git synced 2025-01-05 04:41:03 -05:00
Code Issues Packages Projects Releases Wiki Activity

remove keys from uploadinfo, save media and upload at upload time

Browse Source
This commit is contained in:
Andrea Spacca 2021-01-07 11:22:11 +01:00
parent f0ea2ebd3d
commit 0bf49f7fdd
3 changed files with 37 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

43
pantalaimon/daemon.py
View File

@ -837,7 +837,7 @@ class ProxyDaemon:
except KeyError:
upload_info = self.store.load_upload(self.name, content_uri)
if not upload_info:
return None, None
return None, None, None
self.upload_info[content_uri] = upload_info
@ -845,23 +845,24 @@ class ProxyDaemon:
mxc = urlparse(content_uri)
mxc_server = mxc.netloc.strip("/")
mxc_path = mxc.path.strip("/")
file_name = request.match_info.get("file_name")
logger.info(f"Adding media info for {mxc_server}/{mxc_path} to the store")
media_info = self.store.load_media(self.name, mxc_server, mxc_path)
if not media_info:
return None, None, None
media_info = MediaInfo(mxc_server, mxc_path, upload_info.key, upload_info.iv, upload_info.hashes)
self.media_info[(mxc_server, mxc_path)] = media_info
self.store.save_media(self.name, media_info)
file_name = request.match_info.get("file_name")
return upload_info, media_info, file_name
async def _map_media_upload(self, content_key, content, request, client):
async def _map_decrypted_uri(self, content_key, content, request, client):
try:
upload_info, media_info, file_name = self._get_upload_and_media_info(content_key, content, request)
if not upload_info:
if not upload_info or not media_info:
return await self.forward_to_web(request, token=client.access_token)
response, decrypted_file = await self._load_media(media_info.mcx_server, media_info.mxc_path, file_name)
response, decrypted_file = await self._load_decrypted_file(media_info.mcx_server, media_info.mxc_path, file_name)
if response is None and decrypted_file is None:
return await self.forward_to_web(request, token=client.access_token)
@ -925,7 +926,7 @@ class ProxyDaemon:
content_msgtype = content["msgtype"]
if content_msgtype in ["m.image", "m.video", "m.audio", "m.file"] or msgtype == "m.room.avatar":
try:
content = await self._map_media_upload("url", content, request, client)
content = await self._map_decrypted_uri("url", content, request, client)
return await self.forward_to_web(request, data=json.dumps(content), token=client.access_token)
except ValueError:
return await self.forward_to_web(request, token=client.access_token)
@ -939,7 +940,7 @@ class ProxyDaemon:
content_msgtype = content["msgtype"]
if content_msgtype in ["m.image", "m.video", "m.audio", "m.file"] or msgtype == "m.room.avatar":
upload_info, media_info, file_name = self._get_upload_and_media_info("url", content, request)
if not upload_info:
if not upload_info or not media_info:
response = await client.room_send(
room_id, msgtype, content, txnid, ignore_unverified
)
@ -951,7 +952,11 @@ class ProxyDaemon:
body=await response.transport_response.read(),
)
content = media_info.to_content(file_name, content_msgtype, upload_info.mimetype),
content = media_info.to_content(content["url"],
file_name,
content_msgtype,
upload_info.mimetype
),
response = await client.room_send(
room_id, msgtype, content, txnid, ignore_unverified
@ -1152,7 +1157,15 @@ class ProxyDaemon:
body=await response.transport_response.read(),
)
self.store.save_upload(self.name, response.content_uri, maybe_keys, content_type)
self.store.save_upload(self.name, response.content_uri, content_type)
mxc = urlparse(response.content_uri)
mxc_server = mxc.netloc.strip("/")
mxc_path = mxc.path.strip("/")
logger.info(f"Adding media info for {mxc_server}/{mxc_path} to the store")
media_info = MediaInfo(mxc_server, mxc_path, maybe_keys["key"], maybe_keys["iv"], maybe_keys["hashes"])
self.store.save_media(self.name, media_info)
return web.Response(
status=response.transport_response.status,
@ -1166,7 +1179,7 @@ class ProxyDaemon:
except SendRetryError as e:
return web.Response(status=503, text=str(e))
async def _load_media(self, server_name, media_id, file_name):
async def _load_decrypted_file(self, server_name, media_id, file_name):
try:
media_info = self.media_info[(server_name, media_id)]
except KeyError:
@ -1226,7 +1239,7 @@ class ProxyDaemon:
return self._not_json
try:
content = await self._map_media_upload("avatar_url", content, request, client)
content = await self._map_decrypted_uri("avatar_url", content, request, client)
return await self.forward_to_web(request, data=json.dumps(content), token=client.access_token)
except ValueError:
return await self.forward_to_web(request, token=client.access_token)
@ -1237,7 +1250,7 @@ class ProxyDaemon:
file_name = request.match_info.get("file_name")
try:
response, decrypted_file = await self._load_media(server_name, media_id, file_name)
response, decrypted_file = await self._load_decrypted_file(server_name, media_id, file_name)
if response is None and decrypted_file is None:
return await self.forward_to_web(request)

20
pantalaimon/store.py
View File

@ -48,7 +48,7 @@ class MediaInfo:
iv = attr.ib(type=str)
hashes = attr.ib(type=dict)
def to_content(self, file_name: str, msgtype: str, mime_type: str) -> Dict[Any, Any]:
def to_content(self, url: str, file_name: str, msgtype: str, mime_type: str) -> Dict[Any, Any]:
content = {
"body": file_name,
"file": {
@ -56,7 +56,7 @@ class MediaInfo:
"key": self.key,
"iv": self.iv,
"hashes": self.hashes,
"url": self.url,
"url": url,
"mimetype": mime_type,
}
}
@ -69,9 +69,6 @@ class MediaInfo:
@attr.s
class UploadInfo:
content_uri = attr.ib(type=str)
key = attr.ib(type=dict)
iv = attr.ib(type=str)
hashes = attr.ib(type=dict)
mimetype = attr.ib(type=str)
@ -145,9 +142,7 @@ class PanUploadInfo(Model):
model=Servers, column_name="server_id", backref="upload", on_delete="CASCADE"
)
content_uri = TextField()
key = DictField()
hashes = DictField()
iv = TextField()
mimetype = TextField()
class Meta:
constraints = [SQL("UNIQUE(server_id, content_uri)")]
@ -204,15 +199,12 @@ class PanStore:
return None
@use_database
def save_upload(self, server, content_uri, upload, mimetype):
def save_upload(self, server, content_uri, mimetype):
server = Servers.get(name=server)
PanUploadInfo.insert(
server=server,
content_uri=content_uri,
key=upload["key"],
iv=upload["iv"],
hashes=upload["hashes"],
mimetype=mimetype,
).on_conflict_ignore().execute()
@ -227,7 +219,7 @@ class PanStore:
if i > MAX_LOADED_UPLOAD:
break
upload = UploadInfo(u.content_uri, u.key, u.iv, u.hashes, u.mimetype)
upload = UploadInfo(u.content_uri, u.mimetype)
upload_cache[u.content_uri] = upload
return upload_cache
@ -240,7 +232,7 @@ class PanStore:
if not u:
return None
return UploadInfo(u.content_uri, u.key, u.iv, u.hashes, u.mimetype)
return UploadInfo(u.content_uri, u.mimetype)
@use_database
def save_media(self, server, media):

5
tests/store_test.py
View File

@ -183,13 +183,14 @@ class TestClass(object):
upload_cache = panstore.load_upload(server_name)
assert not upload_cache
mimetype = "image/jpeg"
event = self.encrypted_media_event
assert not panstore.load_upload(server_name, event.url)
upload = UploadInfo(event.url, event.key, event.iv, event.hashes, event.mimetype)
upload = UploadInfo(event.url, mimetype)
panstore.save_upload(server_name, event.url, {"key": event.key, "iv": event.iv, "hashes": event.hashes}, event.mimetype)
panstore.save_upload(server_name, event.url, mimetype)
upload_cache = panstore.load_upload(server_name)