From 465bde5df6f4e8058ca7161075d5b0b2dec12988 Mon Sep 17 00:00:00 2001
From: Ayke van Laethem <aykevanlaethem@gmail.com>
Date: Wed, 27 Aug 2025 09:52:09 +0200
Subject: [PATCH] usb/msc: erase blocks before writing to them

We could perhaps check whether they've been erased already (via unmap or
something) but the default of most OSes seems to be to _not_ unmap
blocks so we will have to erase in most cases anyway.
---
 src/machine/usb/msc/msc.go            | 37 ++++++++++++++++++++++-----
 src/machine/usb/msc/scsi_readwrite.go | 10 ++++++++
 2 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/src/machine/usb/msc/msc.go b/src/machine/usb/msc/msc.go
index d3bf8d6e29..10c8d18972 100644
--- a/src/machine/usb/msc/msc.go
+++ b/src/machine/usb/msc/msc.go
@@ -1,11 +1,13 @@
 package msc
 
 import (
+	"errors"
 	"machine"
 	"machine/usb"
 	"machine/usb/descriptor"
 	"machine/usb/msc/csw"
 	"machine/usb/msc/scsi"
+	"math/bits"
 	"time"
 )
 
@@ -23,6 +25,8 @@ const (
 	mscInterface = 2
 )
 
+var errInvalidBlockSize = errors.New("usb/msc: invalid block size")
+
 var MSC *msc
 
 type msc struct {
@@ -60,20 +64,41 @@ type msc struct {
 }
 
 // Port returns the USB Mass Storage port
-func Port(dev machine.BlockDevice) *msc {
+func Port(dev machine.BlockDevice) (*msc, error) {
 	if MSC == nil {
-		MSC = newMSC(dev)
+		msc, err := newMSC(dev)
+		if err != nil {
+			return nil, err
+		}
+		MSC = msc
 	}
-	return MSC
+	return MSC, nil
 }
 
-func newMSC(dev machine.BlockDevice) *msc {
+func newMSC(dev machine.BlockDevice) (*msc, error) {
 	// Size our buffer to match the maximum packet size of the IN endpoint
 	maxPacketSize := descriptor.EndpointMSCIN.GetMaxPacketSize()
+
+	// Windows only supports block sizes of 512 or 4096 bytes, other systems are
+	// probably similar.
+	blockSize := max(dev.EraseBlockSize(), dev.WriteBlockSize())
+	if bits.OnesCount32(uint32(blockSize)) != 1 {
+		return nil, errInvalidBlockSize // not a power of two
+	}
+	var blockSizeUSB uint32
+	switch {
+	case blockSize <= 512:
+		blockSizeUSB = 512
+	case blockSize <= 4096:
+		blockSizeUSB = 4096
+	default:
+		return nil, errInvalidBlockSize
+	}
+
 	m := &msc{
 		// Some platforms require reads/writes to be aligned to the full underlying hardware block
 		blockCache:    make([]byte, dev.WriteBlockSize()),
-		blockSizeUSB:  512,
+		blockSizeUSB:  blockSizeUSB,
 		buf:           make([]byte, dev.WriteBlockSize()),
 		cswBuf:        make([]byte, csw.MsgLen),
 		cbw:           &CBW{Data: make([]byte, 31)},
@@ -114,7 +139,7 @@ func newMSC(dev machine.BlockDevice) *msc {
 
 	go m.processTasks()
 
-	return m
+	return m, nil
 }
 
 func (m *msc) processTasks() {
diff --git a/src/machine/usb/msc/scsi_readwrite.go b/src/machine/usb/msc/scsi_readwrite.go
index 1b09e13418..5c958a10b4 100644
--- a/src/machine/usb/msc/scsi_readwrite.go
+++ b/src/machine/usb/msc/scsi_readwrite.go
@@ -91,6 +91,16 @@ func (m *msc) writeBlock(b []byte, lba, offset uint32) (n int, err error) {
 		return 0, invalidWriteError
 	}
 
+	// Erase the block first if needed.
+	// Data packets arrive in order, so if we want to write to the start of a
+	// block, that means it's the first write to this erase block and it needs
+	// to be erased first.
+	if uint32(blockStart)%m.blockSizeUSB == 0 {
+		firstBlock := uint32(blockStart) / uint32(m.dev.EraseBlockSize())
+		numBlocks := (m.blockSizeUSB + uint32(m.dev.EraseBlockSize()) - 1) / uint32(m.dev.EraseBlockSize())
+		m.dev.EraseBlocks(int64(firstBlock), int64(numBlocks))
+	}
+
 	// Write the full block to the underlying device
 	n, err = m.dev.WriteAt(b, blockStart)
 	n -= int(blockOffset)