[sw/silicon_creator] Fix short page-misaligned writes in bootstrap

Fixes #14873

This commit also adds a comment highlighting the address wrapping for
the PAGE_PROGRAM command.

Signed-off-by: Alphan Ulusoy <alphan@google.com>
diff --git a/sw/device/silicon_creator/rom/bootstrap.c b/sw/device/silicon_creator/rom/bootstrap.c
index 0ec2c6e..e7ee645 100644
--- a/sw/device/silicon_creator/rom/bootstrap.c
+++ b/sw/device/silicon_creator/rom/bootstrap.c
@@ -179,15 +179,21 @@
       .write = kMultiBitBool4True,
       .erase = kMultiBitBool4False,
   });
-  // Wrap to the start of the flash programming page if `addr` is not aligned.
+  // Perform two writes if the start address is not page-aligned (256 bytes).
+  // Note: Address is flash-word-aligned (8 bytes) due to the check above.
   rom_error_t err_0 = kErrorOk;
   size_t prog_page_misalignment = addr & kFlashProgPageMask;
   if (prog_page_misalignment > 0) {
     size_t word_count =
         (kFlashProgPageSize - prog_page_misalignment) / sizeof(uint32_t);
+    if (word_count > rem_word_count) {
+      word_count = rem_word_count;
+    }
     err_0 = flash_ctrl_data_write(addr, word_count, data);
     rem_word_count -= word_count;
     data += word_count * sizeof(uint32_t);
+    // Wrap to the beginning of the current page since PAGE_PROGRAM modifies
+    // a single page only.
     addr &= ~kFlashProgPageMask;
   }
   rom_error_t err_1 = kErrorOk;
diff --git a/sw/device/silicon_creator/rom/bootstrap_unittest.cc b/sw/device/silicon_creator/rom/bootstrap_unittest.cc
index 1f3f987..926a9e6 100644
--- a/sw/device/silicon_creator/rom/bootstrap_unittest.cc
+++ b/sw/device/silicon_creator/rom/bootstrap_unittest.cc
@@ -327,7 +327,7 @@
   EXPECT_EQ(bootstrap(), kErrorUnknown);
 }
 
-TEST_F(BootstrapTest, BootstrapMisalignedAddr) {
+TEST_F(BootstrapTest, BootstrapMisalignedAddrLongPayload) {
   // Erase
   ExpectBootstrapRequestCheck(true);
   EXPECT_CALL(spi_device_, Init());
@@ -361,6 +361,37 @@
   EXPECT_EQ(bootstrap(), kErrorUnknown);
 }
 
+TEST_F(BootstrapTest, BootstrapMisalignedAddrShortPayload) {
+  // Erase
+  ExpectBootstrapRequestCheck(true);
+  EXPECT_CALL(spi_device_, Init());
+  ExpectSpiCmd(ChipEraseCmd());
+  ExpectSpiFlashStatusGet(true);
+  ExpectFlashCtrlChipErase(kErrorOk, kErrorOk);
+  // Verify
+  ExpectFlashCtrlEraseVerify(kErrorOk, kErrorOk);
+  EXPECT_CALL(spi_device_, FlashStatusClear());
+  // Program
+  auto cmd = PageProgramCmd(816, 8);
+  ExpectSpiCmd(cmd);
+  ExpectSpiFlashStatusGet(true);
+
+  std::vector<uint8_t> flash_bytes(cmd.payload,
+                                   cmd.payload + cmd.payload_byte_count);
+
+  ExpectFlashCtrlWriteEnable();
+  EXPECT_CALL(flash_ctrl_, DataWrite(816, 2, HasBytes(flash_bytes)))
+      .WillOnce(Return(kErrorOk));
+  ExpectFlashCtrlAllDisable();
+
+  EXPECT_CALL(spi_device_, FlashStatusClear());
+  // Reset
+  ExpectSpiCmd(ResetCmd());
+  EXPECT_CALL(rstmgr_, Reset());
+
+  EXPECT_EQ(bootstrap(), kErrorUnknown);
+}
+
 TEST_F(BootstrapTest, BootstrapStartWithSectorErase) {
   // Erase
   ExpectBootstrapRequestCheck(true);