diff --git a/frontend/src/StartView.svelte b/frontend/src/StartView.svelte index a97607eed..3a8358e98 100644 --- a/frontend/src/StartView.svelte +++ b/frontend/src/StartView.svelte @@ -121,8 +121,21 @@ tryHash = !!window.location.hash; let mouseX, selectedAnimal; const warnFileSize = 250 * 1024 * 1024; + const fileChunkSize = warnFileSize; + + async function sendChunk(id, f, start) { + let end = Math.min(start + fileChunkSize, f.size); + await fetch( + `${$settings.backendUrl}/root_resource_chunk?id=${id}&start=${start}&end=${end}`, + { + method: "POST", + body: await f.slice(start, end), + } + ); + } async function createRootResource(f) { + let rootModel; if ( f.size > warnFileSize && !window.confirm( @@ -134,15 +147,34 @@ showRootResource = false; return; } + if (f.size > warnFileSize) { + let id = await fetch( + `${$settings.backendUrl}/init_chunked_root_resource?name=${f.name}&size=${f.size}`, + { method: "POST" } + ).then((r) => r.json()); + let chunkStartAddrs = Array.from( + { length: Math.ceil(f.size / fileChunkSize) }, + (v, i) => i * fileChunkSize + ); + await Promise.all( + chunkStartAddrs.map((start) => sendChunk(id, f, start)) + ); - const rootModel = await fetch( - `${$settings.backendUrl}/create_root_resource?name=${f.name}`, - { - method: "POST", - body: await f.arrayBuffer(), - } - ).then((r) => r.json()); - + rootModel = await fetch( + `${$settings.backendUrl}/create_chunked_root_resource?id=${id}&name=${f.name}`, + { + method: "POST", + } + ).then((r) => r.json()); + } else { + rootModel = await fetch( + `${$settings.backendUrl}/create_root_resource?name=${f.name}`, + { + method: "POST", + body: await f.arrayBuffer(), + } + ).then((r) => r.json()); + } rootResource = remote_model_to_resource(rootModel, resources); $selected = rootModel.id; } diff --git a/ofrak_core/CHANGELOG.md b/ofrak_core/CHANGELOG.md index d0bcc7a23..e17300677 100644 --- a/ofrak_core/CHANGELOG.md +++ b/ofrak_core/CHANGELOG.md @@ -8,6 +8,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/) - Add a JFFS2 packer and unpacker. ([#326](https://github.com/redballoonsecurity/ofrak/pull/326)) ### Changed +- Support uploading files in chunks to handle files larger than 2GB from the GUI ([#324](https://github.com/redballoonsecurity/ofrak/pull/324)) ### Fixed diff --git a/ofrak_core/ofrak/gui/server.py b/ofrak_core/ofrak/gui/server.py index 870440764..cd521c057 100644 --- a/ofrak_core/ofrak/gui/server.py +++ b/ofrak_core/ofrak/gui/server.py @@ -144,9 +144,13 @@ def __init__( self.resource_view_context: ResourceViewContext = ResourceViewContext() self.component_context: ComponentContext = ClientComponentContext() self.script_builder: ScriptBuilder = ScriptBuilder() + self.resource_builder: Dict[str, Tuple[Resource, memoryview]] = {} self._app.add_routes( [ web.post("/create_root_resource", self.create_root_resource), + web.post("/init_chunked_root_resource", self.init_chunked_root_resource), + web.post("/root_resource_chunk", self.root_resource_chunk), + web.post("/create_chunked_root_resource", self.create_chunked_root_resource), web.get("/get_root_resources", self.get_root_resources), web.get("/{resource_id}/", self.get_resource), web.get("/{resource_id}/get_data", self.get_data), @@ -260,6 +264,66 @@ async def run_until_cancelled(self): # pragma: no cover finally: await self.runner.cleanup() + @exceptions_to_http(SerializedError) + async def init_chunked_root_resource(self, request: Request) -> Response: + name = request.query.get("name") + size_param = request.query.get("size") + if name is None: + raise HTTPBadRequest(reason="Missing resource name from request") + if size_param is None: + raise HTTPBadRequest(reason="Missing chunk size from request") + size = int(size_param) + root_resource: Resource = await self._ofrak_context.create_root_resource(name, b"", (File,)) + self.resource_builder[root_resource.get_id().hex()] = ( + root_resource, + memoryview(bytearray(b"\x00" * size)), + ) + return json_response(root_resource.get_id().hex()) + + @exceptions_to_http(SerializedError) + async def root_resource_chunk(self, request: Request) -> Response: + id = request.query.get("id") + start_param = request.query.get("start") + end_param = request.query.get("end") + if id is None: + raise HTTPBadRequest(reason="Missing resource id from request") + if start_param is None: + raise HTTPBadRequest(reason="Missing chunk start from request") + if end_param is None: + raise HTTPBadRequest(reason="Missing chunk end from request") + start = int(start_param) + end = int(end_param) + chunk_data = await request.read() + _, data = self.resource_builder[id] + data[start:end] = chunk_data + return json_response([]) + + @exceptions_to_http(SerializedError) + async def create_chunked_root_resource(self, request: Request) -> Response: + id = request.query.get("id") + name = request.query.get("name") + if id is None: + return HTTPBadRequest(reason="Missing root resource `id` from request") + if name is None: + return HTTPBadRequest(reason="Missing root resource `name` from request") + + try: + root_resource, data = self.resource_builder[id] + script_str = rf""" + if root_resource is None: + root_resource = await ofrak_context.create_root_resource_from_file("{name}")""" + root_resource.queue_patch(Range(0, 0), bytearray(data)) + await root_resource.save() + await self.script_builder.add_action(root_resource, script_str, ActionType.UNPACK) + if request.remote is not None: + self._job_ids[request.remote] = root_resource.get_job_id() + await self.script_builder.commit_to_script(root_resource) + except Exception as e: + await self.script_builder.clear_script_queue(root_resource) + raise e + self.resource_builder.pop(id) + return json_response(self._serialize_resource(root_resource)) + @exceptions_to_http(SerializedError) async def create_root_resource(self, request: Request) -> Response: name = request.query.get("name") diff --git a/ofrak_core/test_ofrak/unit/test_ofrak_server.py b/ofrak_core/test_ofrak/unit/test_ofrak_server.py index 09af5b414..28da27aca 100644 --- a/ofrak_core/test_ofrak/unit/test_ofrak_server.py +++ b/ofrak_core/test_ofrak/unit/test_ofrak_server.py @@ -1,6 +1,9 @@ import itertools import json import os +import tempfile +from ofrak.ofrak_context import OFRAKContext +from ofrak.resource import Resource import pytest import re import sys @@ -26,6 +29,14 @@ def hello_world_elf() -> bytes: return hello_elf() +@pytest.fixture() +async def large_test_file(ofrak_context: OFRAKContext) -> Resource: + with tempfile.NamedTemporaryFile() as temp: + for i in range(256): + temp.write(int.to_bytes(i, 1, "big") * 1024 * 1024) + yield await ofrak_context.create_root_resource_from_file(temp.name) + + @pytest.fixture(scope="session") def firmware_zip() -> bytes: assets_dir = os.path.abspath( @@ -119,6 +130,32 @@ async def test_create_root_resource( assert body["tags"] == json_result["tags"] +async def test_create_chunked_root_resource( + ofrak_client: TestClient, ofrak_server, large_test_file +): + test_file_data = await large_test_file.get_data() + chunk_size = int(len(test_file_data) / 10) + init_resp = await ofrak_client.post( + "/init_chunked_root_resource", + params={"name": "test_file_data", "size": len(test_file_data)}, + ) + id = await init_resp.json() + for start in range(0, len(test_file_data), chunk_size): + end = min(start + chunk_size, len(test_file_data)) + res = await ofrak_client.post( + "/root_resource_chunk", + params={"id": id, "start": start, "end": end}, + data=test_file_data[start:end], + ) + create_resp = await ofrak_client.post( + "/create_chunked_root_resource", params={"name": "test_file_data", "id": id} + ) + assert create_resp.status == 200 + length_resp = await ofrak_client.get(f"/{id}/get_data_length") + length_resp_body = await length_resp.json() + assert length_resp_body == len(test_file_data) + + async def test_get_root_resources( ofrak_client: TestClient, ofrak_context, ofrak_server, hello_world_elf ):