From 5bcb3dba3fcb259d8d191e68aff27e57410a3960 Mon Sep 17 00:00:00 2001
From: "Brian S. Stephan" <bss@incorporeal.org>
Date: Sun, 9 Jul 2023 09:41:35 -0500
Subject: [PATCH] method to write the config over USB at the proper alignment

---
 gp2040ce_bintools/builder.py | 21 +++++++++++++++++++--
 tests/test_builder.py        | 20 ++++++++++++++++++--
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/gp2040ce_bintools/builder.py b/gp2040ce_bintools/builder.py
index b660649..3c51d5c 100644
--- a/gp2040ce_bintools/builder.py
+++ b/gp2040ce_bintools/builder.py
@@ -6,8 +6,9 @@ import logging
 from google.protobuf.message import Message
 
 from gp2040ce_bintools import core_parser
-from gp2040ce_bintools.storage import (STORAGE_BINARY_LOCATION, STORAGE_SIZE, pad_config_to_storage_size,
-                                       serialize_config_with_footer)
+from gp2040ce_bintools.pico import write
+from gp2040ce_bintools.storage import (STORAGE_BINARY_LOCATION, STORAGE_MEMORY_ADDRESS, STORAGE_SIZE,
+                                       pad_config_to_storage_size, serialize_config_with_footer)
 
 logger = logging.getLogger(__name__)
 
@@ -113,6 +114,22 @@ def write_new_config_to_filename(config: Message, filename: str, inject: bool =
         file.write(binary)
 
 
+def write_new_config_to_usb(config: Message, endpoint_out: object, endpoint_in: object):
+    """Serialize the provided config to a device over USB, in the proper location for a GP2040-CE board.
+
+    Args:
+        config: the Protobuf configuration to write to a Pico board in BOOTSEL mode
+        endpoint_out: the USB endpoint to write to
+        endpoint_in: the USB endpoint to read from
+    """
+    serialized = serialize_config_with_footer(config)
+    # we don't write the whole area, just the minimum from the end of the storage section
+    # nevertheless, the USB device needs writes to start at 256 byte boundaries
+    padding = 256 - (len(serialized) % 256)
+    binary = bytearray(b'\x00' * padding) + serialized
+    write(endpoint_out, endpoint_in, STORAGE_MEMORY_ADDRESS + (STORAGE_SIZE - len(binary)), binary)
+
+
 ############
 # COMMANDS #
 ############
diff --git a/tests/test_builder.py b/tests/test_builder.py
index f1f8484..b5e0fc4 100644
--- a/tests/test_builder.py
+++ b/tests/test_builder.py
@@ -1,14 +1,15 @@
 """Tests for the image builder module."""
 import os
 import sys
+import unittest.mock as mock
 
 import pytest
 from decorator import decorator
 
 from gp2040ce_bintools import get_config_pb2
 from gp2040ce_bintools.builder import (FirmwareLengthError, combine_firmware_and_config, pad_firmware_up_to_storage,
-                                       replace_config_in_binary, write_new_config_to_filename)
-from gp2040ce_bintools.storage import get_config, get_config_footer, get_storage_section
+                                       replace_config_in_binary, write_new_config_to_filename, write_new_config_to_usb)
+from gp2040ce_bintools.storage import get_config, get_config_footer, get_storage_section, serialize_config_with_footer
 
 HERE = os.path.dirname(os.path.abspath(__file__))
 
@@ -134,3 +135,18 @@ def test_write_new_config_to_config_bin(firmware_binary, tmp_path):
     config_size, _, _ = get_config_footer(config_dump)
     assert config.boardVersion == 'v0.7.2-COOL'
     assert len(config_dump) == config_size + 12
+
+
+@with_pb2s
+def test_write_new_config_to_usb(config_binary):
+    """Test that the config can be written to USB at the proper alignment."""
+    config = get_config(config_binary)
+    serialized = serialize_config_with_footer(config)
+    end_out, end_in = mock.MagicMock(), mock.MagicMock()
+    with mock.patch('gp2040ce_bintools.builder.write') as mock_write:
+        write_new_config_to_usb(config, end_out, end_in)
+
+    # check that it got padded
+    padded_serialized = bytearray(b'\x00' * 4) + serialized
+    assert mock_write.call_args.args[2] % 256 == 0
+    assert mock_write.call_args.args[3] == padded_serialized