From 0ae2d82451dc8fc1f66cc0e144e40a6936af93d1 Mon Sep 17 00:00:00 2001
From: teri-anric <2005ahi2005@gmail.com>
Date: Sat, 17 Aug 2024 23:25:32 +0300
Subject: [PATCH] Fix readable BufferedInputFile and add support BufferedIOBase

---
 CHANGES/1564.misc.rst                        |  5 ++++
 aiogram/types/input_file.py                  | 24 ++++++++++++++------
 tests/test_api/test_types/test_input_file.py | 21 +++++++++++++++--
 3 files changed, 41 insertions(+), 9 deletions(-)
 create mode 100644 CHANGES/1564.misc.rst

diff --git a/CHANGES/1564.misc.rst b/CHANGES/1564.misc.rst
new file mode 100644
index 00000000..ecaed001
--- /dev/null
+++ b/CHANGES/1564.misc.rst
@@ -0,0 +1,5 @@
+Fixed unsafe file reading in BufferedInputFile.from_file and 
+added support for initializing BufferedInputFile from BytesIO and 
+other bytes buffers. 
+This change improves safety and extends flexibility in handling 
+different byte sources when working with buffered input files.
diff --git a/aiogram/types/input_file.py b/aiogram/types/input_file.py
index 5b730598..dfbb5179 100644
--- a/aiogram/types/input_file.py
+++ b/aiogram/types/input_file.py
@@ -38,7 +38,12 @@ class InputFile(ABC):
 
 
 class BufferedInputFile(InputFile):
-    def __init__(self, file: bytes, filename: str, chunk_size: int = DEFAULT_CHUNK_SIZE):
+    def __init__(
+        self,
+        buffer_or_bytes: bytes | io.BufferedIOBase,
+        filename: str,
+        chunk_size: int = DEFAULT_CHUNK_SIZE,
+    ):
         """
         Represents object for uploading files from filesystem
 
@@ -48,7 +53,10 @@ class BufferedInputFile(InputFile):
         """
         super().__init__(filename=filename, chunk_size=chunk_size)
 
-        self.data = file
+        if isinstance(buffer_or_bytes, bytes):
+            buffer_or_bytes = io.BytesIO(buffer_or_bytes)
+
+        self.buffer = buffer_or_bytes
 
     @classmethod
     def from_file(
@@ -68,13 +76,15 @@ class BufferedInputFile(InputFile):
         """
         if filename is None:
             filename = os.path.basename(path)
-        with open(path, "rb") as f:
-            data = f.read()
-        return cls(data, filename=filename, chunk_size=chunk_size)
+        buffer = io.BytesIO()
+        with open(path, "rb") as fp:
+            while chunk := fp.read(chunk_size):
+                buffer.write(chunk)
+        return cls(buffer, filename=filename, chunk_size=chunk_size)
 
     async def read(self, bot: "Bot") -> AsyncGenerator[bytes, None]:
-        buffer = io.BytesIO(self.data)
-        while chunk := buffer.read(self.chunk_size):
+        self.buffer.seek(0)
+        while chunk := self.buffer.read(self.chunk_size):
             yield chunk
 
 
diff --git a/tests/test_api/test_types/test_input_file.py b/tests/test_api/test_types/test_input_file.py
index e8716f84..4b89fbb8 100644
--- a/tests/test_api/test_types/test_input_file.py
+++ b/tests/test_api/test_types/test_input_file.py
@@ -1,3 +1,4 @@
+import io
 from typing import AsyncIterable
 
 from aresponses import ResponsesMockServer
@@ -34,7 +35,7 @@ class TestInputFile:
 
         assert isinstance(file, InputFile)
         assert file.filename == "file.bin"
-        assert isinstance(file.data, bytes)
+        assert isinstance(file.buffer, io.BytesIO)
 
     async def test_buffered_input_file_readable(self, bot: MockedBot):
         file = BufferedInputFile(b"\f" * 10, filename="file.bin", chunk_size=1)
@@ -53,7 +54,7 @@ class TestInputFile:
         assert file.filename is not None
         assert file.filename.startswith("test_")
         assert file.filename.endswith(".py")
-        assert isinstance(file.data, bytes)
+        assert isinstance(file.buffer, io.BytesIO)
         assert file.chunk_size == 10
 
     async def test_buffered_input_file_from_file_readable(self, bot: MockedBot):
@@ -66,6 +67,22 @@ class TestInputFile:
             size += chunk_size
         assert size > 0
 
+    async def test_buffered_input_file_from_buffer(self, bot: MockedBot):
+        buffer = io.BytesIO(b"\f" * 10)
+        file = BufferedInputFile(buffer, filename="file.bin", chunk_size=1)
+
+        assert file.buffer is buffer
+        assert file.filename == "file.bin"
+
+        size = 0
+        async for chunk in file.read(bot):
+            chunk_size = len(chunk)
+            assert isinstance(chunk, bytes)
+            assert chunk_size == 1
+            size += chunk_size
+        assert size == 10
+        assert file.buffer.seek(0, io.SEEK_CUR) == 10
+
     async def test_url_input_file(self, aresponses: ResponsesMockServer):
         aresponses.add(
             aresponses.ANY,