fix: compute SHA-256 digest at correct image offset for combined binaries

pkg/espflasher/flasher.go

@@ -357,9 +357,12 @@ func (f *Flasher) FlashImage(data []byte, offset uint32, progress ProgressFunc)
 		}
 	}
 
-	// Use compressed flash only if supported (ESP8266 ROM doesn't support it)
-	canCompress := f.chip == nil || f.chip.ROMHasCompressedFlash || f.conn.isStub()
-	if f.opts.Compress && canCompress {
+	// Use compressed flash only with the stub loader. While most ESP32+ ROM
+	// bootloaders support compressed flash commands, we must skip flashDeflEnd
+	// for ROM (it exits the bootloader), which can leave data unflushed in
+	// the ROM's decompressor buffer. esptool also defaults to uncompressed
+	// writes for ROM (compress = IS_STUB).
+	if f.opts.Compress && f.conn.isStub() {
 		return f.flashCompressed(data, offset, progress)
 	}
 	return f.flashUncompressed(data, offset, progress)
@@ -408,9 +411,6 @@ func (f *Flasher) FlashImages(images []ImagePart, progress ProgressFunc) error {
 		}
 	}
 
-	// Determine if compressed flash is available
-	canCompress := f.chip == nil || f.chip.ROMHasCompressedFlash || f.conn.isStub()
-
 	totalSize := 0
 	for _, img := range images {
 		totalSize += len(img.Data)
@@ -442,7 +442,7 @@ func (f *Flasher) FlashImages(images []ImagePart, progress ProgressFunc) error {
 			}
 		}
 
-		if f.opts.Compress && canCompress {
+		if f.opts.Compress && f.conn.isStub() {
 			err = f.flashCompressed(data, img.Offset, partProgress)
 		} else {
 			err = f.flashUncompressed(data, img.Offset, partProgress)
@@ -515,9 +515,18 @@ func (f *Flasher) flashCompressed(data []byte, offset uint32, progress ProgressF
 		}
 	}
 
-	// End flash
-	if err := f.conn.flashDeflEnd(false); err != nil {
-		return err
+	// End the compressed flash session.
+	// For ROM bootloaders, skip sending FLASH_DEFL_END — the ROM exits the
+	// bootloader upon receiving it, which can interfere with flash operations.
+	// esptool also skips this for ROM: "skip sending flash_finish to ROM loader,
+	// as it causes the loader to exit and run user code."
+	// For the stub, the end command acts as a write barrier: the stub ACKs each
+	// block on receive but writes to flash asynchronously, so the end command
+	// ensures the last block is actually written before we proceed.
+	if f.conn.isStub() {
+		if err := f.conn.flashDeflEnd(false); err != nil {
+			return err
+		}
 	}
 
 	f.logf("Flash complete. Verifying...")
@@ -577,9 +586,14 @@ func (f *Flasher) flashUncompressed(data []byte, offset uint32, progress Progres
 		}
 	}
 
-	// End flash
-	if err := f.conn.flashEnd(false); err != nil {
-		return err
+	// End the flash session.
+	// For ROM bootloaders, skip sending FLASH_END — the ROM exits the
+	// bootloader upon receiving it. esptool also skips this for ROM.
+	// For the stub, the end command acts as a write barrier.
+	if f.conn.isStub() {
+		if err := f.conn.flashEnd(false); err != nil {
+			return err
+		}
 	}
 
 	f.logf("Flash complete. Verifying...")
@@ -656,11 +670,18 @@ func (f *Flasher) WriteRegister(addr, value uint32) error {
 
 // Reset performs a hard reset of the device, causing it to run user code.
 func (f *Flasher) Reset() {
-	// Try to cleanly exit the stub/rom loader
-	f.conn.flashBegin(0, 0, false) //nolint:errcheck
-	f.conn.flashEnd(true)          //nolint:errcheck
-	time.Sleep(50 * time.Millisecond)
+	if f.conn.isStub() {
+		// The stub loader needs an explicit flash_begin/flash_end to
+		// cleanly exit flash mode before hardware reset.
+		f.conn.flashBegin(0, 0, false) //nolint:errcheck
+		f.conn.flashEnd(true)          //nolint:errcheck
+		time.Sleep(50 * time.Millisecond)
+	}
 
+	// For ROM bootloaders, skip flash_begin/flash_end — sending
+	// CMD_FLASH_BEGIN after a compressed download may interfere with
+	// the flash controller state at offset 0. esptool also just does
+	// a hard reset without any flash commands for the ROM path.
 	hardReset(f.port, false)
 	f.logf("Device reset.")
 }

pkg/espflasher/image.go

@@ -2,6 +2,7 @@ package espflasher
 
 import (
 	"crypto/sha256"
+	"encoding/binary"
 	"fmt"
 )
 
@@ -108,15 +109,67 @@ func (f *Flasher) patchImageHeader(data []byte) ([]byte, error) {
 
 	// Update SHA256 hash if present.
 	// The extended header is at bytes 8-23, and byte 23 bit 0 indicates SHA256
-	// is appended as the last 32 bytes. ESP8266 doesn't have an extended header,
-	// so skip SHA256 update for it.
+	// is appended after the image content. ESP8266 doesn't have an extended
+	// header, so skip SHA256 update for it.
+	//
+	// We parse the image structure to find the exact offset of the SHA256
+	// digest rather than assuming it is the last 32 bytes. This is critical
+	// for combined/merged binaries where the bootloader image is followed by
+	// partition table and application data.
 	if f.chip != nil && f.chip.ChipType != ChipESP8266 &&
 		len(patched) >= 24+32 && patched[23]&1 != 0 {
-		content := patched[:len(patched)-32]
-		hash := sha256.Sum256(content)
-		copy(patched[len(patched)-32:], hash[:])
-		f.logf("SHA digest in image updated")
+		dataLen := imageDataLength(patched)
+		if dataLen >= 0 && dataLen+32 <= len(patched) {
+			content := patched[:dataLen]
+			hash := sha256.Sum256(content)
+			copy(patched[dataLen:dataLen+32], hash[:])
+			f.logf("SHA digest in image updated")
+		}
 	}
 
 	return patched, nil
 }
+
+// imageDataLength parses an ESP32 firmware image to determine the byte length
+// of the image content before the appended SHA256 digest.
+//
+// The image structure is:
+//   - Common header: 8 bytes (magic, segment_count, flash_mode, size_freq, entry_point)
+//   - Extended header: 16 bytes (wp_pin, drive levels, chip_id, revisions, append_digest)
+//   - N segments: each has an 8-byte header (addr, length) followed by data bytes
+//   - Padding to align to 16 bytes + 1-byte checksum
+//   - SHA256 digest: 32 bytes (if append_digest is set)
+//
+// Returns the offset where the SHA256 digest starts, or -1 if the image
+// cannot be parsed. This is equivalent to esptool's data_length field.
+func imageDataLength(data []byte) int {
+	if len(data) < 24 || data[0] != espImageMagic {
+		return -1
+	}
+
+	segCount := int(data[1])
+	pos := 24 // common header (8) + extended header (16)
+
+	for i := 0; i < segCount; i++ {
+		if pos+8 > len(data) {
+			return -1
+		}
+		segLen := int(binary.LittleEndian.Uint32(data[pos+4 : pos+8]))
+		pos += 8 + segLen
+		if pos > len(data) {
+			return -1
+		}
+	}
+
+	// The checksum byte is placed so the file position becomes a multiple of 16.
+	// This matches esptool's align_file_position(f, 16) + read(1):
+	//   align = (16 - 1) - (pos % 16)   →  skip 'align' padding bytes
+	//   read 1 byte (checksum)
+	//   final position = pos + 16 - (pos % 16)
+	dataLen := pos + 16 - (pos % 16)
+	if dataLen > len(data) {
+		return -1
+	}
+
+	return dataLen
+}