Support unknown file sizes in streaming encryption to fix SFTP timeouts

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:

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:

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)

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