diff --git a/dcfs/crypto/cipher.py b/dcfs/crypto/cipher.py index d3a8aef..911201c 100644 --- a/dcfs/crypto/cipher.py +++ b/dcfs/crypto/cipher.py @@ -134,6 +134,8 @@ def ciphertext_size_for_plaintext(plaintext_size: int, chunk_size: int) -> int: must be added separately. Used by the upload wrapper to declare the correct file size up-front to Discord. """ + if plaintext_size == -1: + return -1 if plaintext_size < 0: raise ValueError(f"plaintext size negative: {plaintext_size}") if plaintext_size == 0: diff --git a/dcfs/crypto/stream.py b/dcfs/crypto/stream.py index b5483a0..d92636f 100644 --- a/dcfs/crypto/stream.py +++ b/dcfs/crypto/stream.py @@ -87,9 +87,10 @@ def wrap( raise RuntimeError("serialized header has unexpected size") plaintext_size = inner.get_size() - ciphertext_size = HEADER_SIZE + ciphertext_size_for_plaintext( + ct_payload_size = ciphertext_size_for_plaintext( plaintext_size, header.chunk_size ) + ciphertext_size = -1 if ct_payload_size == -1 else HEADER_SIZE + ct_payload_size msg = cls( name=inner.name, @@ -165,10 +166,11 @@ async def read(self, length: int) -> bytes: if not self._header_emitted: await self.open() - remaining = self.size - self._read_size - if remaining <= 0: - return b"" - length = min(length, remaining) + if self.size >= 0: + remaining = self.size - self._read_size + if remaining <= 0: + return b"" + length = min(length, remaining) out = bytearray() while len(out) < length: @@ -204,11 +206,16 @@ async def _read_plaintext_chunk(self) -> bytes: """ target = self.cipher.chunk_size buf = bytearray() - remaining_plaintext = self.inner.get_size() - self._plaintext_read - if remaining_plaintext <= 0: - return b"" - want = min(target, remaining_plaintext) + plaintext_total = self.inner.get_size() + if plaintext_total >= 0: + remaining_plaintext = plaintext_total - self._plaintext_read + if remaining_plaintext <= 0: + return b"" + want = min(target, remaining_plaintext) + else: + want = target + while len(buf) < want: piece = await self.inner.read(want - len(buf)) if not piece: diff --git a/tests/test_crypto/test_cipher.py b/tests/test_crypto/test_cipher.py index 32e55d7..9fae356 100644 --- a/tests/test_crypto/test_cipher.py +++ b/tests/test_crypto/test_cipher.py @@ -142,6 +142,9 @@ def test_ciphertext_size_for_plaintext( ciphertext_size_for_plaintext(plaintext_size, chunk_size) == expected ) + def test_ciphertext_size_for_unknown_plaintext_size(self) -> None: + assert ciphertext_size_for_plaintext(-1, 4096) == -1 + def test_offset_to_chunk_round_trip(self) -> None: for offset in (0, 1, 64 * 1024 - 1, 64 * 1024, 64 * 1024 + 5): chunk_index, in_chunk = plaintext_offset_to_chunk(offset, 64 * 1024) diff --git a/tests/test_crypto/test_repository.py b/tests/test_crypto/test_repository.py index 1413997..4bc6234 100644 --- a/tests/test_crypto/test_repository.py +++ b/tests/test_crypto/test_repository.py @@ -26,6 +26,7 @@ from dcfs.reqres import ( FileContent, FileMessageFromBuffer, + FileMessageFromStream, SentFileMessage, UploadableFileMessage, ) @@ -69,15 +70,16 @@ async def save( await file_msg.open() try: buf = bytearray() + total_size = file_msg.get_size() for read_size in (7, 4096, 8192, 1024 * 1024): while True: piece = await file_msg.read(read_size) if not piece: break buf += piece - if len(buf) >= file_msg.get_size(): + if total_size >= 0 and len(buf) >= total_size: break - if len(buf) >= file_msg.get_size(): + if not buf or (total_size >= 0 and len(buf) >= total_size): break finally: await file_msg.close() @@ -399,3 +401,27 @@ async def test_detection_cached_across_calls() -> None: hit, entry = repo._cache.lookup(fv.id) assert hit and entry is None # cached as plaintext + + +async def test_streaming_upload_unknown_size_round_trip() -> None: + repo = _make_repo(chunk_size=4096) + plaintext = os.urandom(4096 * 2 + 100) + + async def gen(): + yield plaintext[:4096] + yield plaintext[4096:8192] + yield plaintext[8192:] + + # size=-1 indicates unknown size + file_msg = FileMessageFromStream.new(stream=gen(), size=-1, name="stream.bin") + sent = await repo.save(file_msg) + + fv = _FakeFileVersion( + id="v-stream", + size=sum(m.size for m in sent), + message_ids=[m.message_id for m in sent], + part_sizes=[m.size for m in sent], + ) + + out = await _collect(await repo.get(fv, 0, -1, "stream.bin")) + assert out == plaintext