diff --git a/pkg/espflasher/flasher.go b/pkg/espflasher/flasher.go index b4d7fa4..e3014f2 100644 --- a/pkg/espflasher/flasher.go +++ b/pkg/espflasher/flasher.go @@ -81,9 +81,10 @@ type Logger interface { // ProgressFunc is called with progress updates during long-running // operations. The meaning of current and total is operation-defined: for // flashing and reading, they are bytes transferred and total bytes; for -// erase, they are elapsed and estimated-total milliseconds, since erase is -// a single blocking command with no per-chunk protocol seam to report exact -// byte progress. +// erase and GetFlashMD5, they are elapsed and estimated-total milliseconds, +// since both are single blocking commands with no per-chunk protocol seam +// to report exact byte progress — the device stays silent for the duration +// and reports no real completion percentage. type ProgressFunc func(current, total int) // ConnectPhase identifies a stage of the bootloader connection sequence. @@ -985,7 +986,13 @@ func (f *Flasher) GetSecurityInfo() (*SecurityInfo, error) { // GetFlashMD5 returns the MD5 hash of a flash region. // Requires the stub loader to be running. -func (f *Flasher) GetFlashMD5(offset, size uint32) (string, error) { +// +// If progress is non-nil, it is called periodically with a synthetic ETA +// (elapsed and estimated milliseconds) ticked against the same size-scaled +// timeout used internally for the MD5 command, since the device stays +// silent for the duration of the hash computation and reports no real +// completion percentage. Pass nil to skip this entirely. +func (f *Flasher) GetFlashMD5(offset, size uint32, progress ProgressFunc) (string, error) { if !f.conn.isStub() { return "", &UnsupportedCommandError{Command: "flash MD5 (requires stub)"} } @@ -994,13 +1001,103 @@ func (f *Flasher) GetFlashMD5(offset, size uint32) (string, error) { return "", err } - result, err := f.conn.flashMD5(offset, size) + if progress == nil { + result, err := f.conn.flashMD5(offset, size) + if err != nil { + return "", err + } + return hex.EncodeToString(result), nil + } + + var result []byte + err := tickMD5(md5TimeoutForSize(size), md5ProgressInterval, progress, func() error { + r, err := f.conn.flashMD5(offset, size) + if err != nil { + return err + } + result = r + return nil + }) if err != nil { return "", err } return hex.EncodeToString(result), nil } +// md5ProgressInterval is the tick interval used by tickMD5 when reporting +// synthetic MD5 progress via GetFlashMD5. +const md5ProgressInterval = 500 * time.Millisecond + +// tickMD5 runs work (a blocking flash-MD5 call) while emitting synthetic ETA +// progress updates against est every interval, until work returns. progress +// is called with (elapsedMs, estMs), capped so elapsed never reaches est +// until work has actually completed successfully. On success, a final +// progress(estMs, estMs) call is emitted; on error, no such call is made. +// +// Intermediate ETA ticks are best-effort: they run on a background ticker +// goroutine, so a panic from the caller's progress callback during such a +// tick is recovered and that tick is dropped silently rather than crashing +// the process. The final completion progress(estMs, estMs) call runs +// synchronously in the caller's goroutine (via the deferred cleanup below), +// so a panic there propagates to the caller normally. +// +// The ticker goroutine only ever calls progress — it never touches the +// connection — and is always stopped and joined before tickMD5 returns, via +// a deferred cleanup registered before work is invoked. This holds even if +// work panics: the deferred cleanup still runs during unwinding, the ticker +// goroutine is joined, and the panic is left to propagate afterward without +// emitting a bogus final progress(estMs, estMs) call. +func tickMD5(est, interval time.Duration, progress ProgressFunc, work func() error) error { + estMs := int(est / time.Millisecond) + if estMs < 1 { + estMs = 1 + } + + done := make(chan struct{}) + stopped := make(chan struct{}) + + go func() { + defer close(stopped) + ticker := time.NewTicker(interval) + defer ticker.Stop() + + start := time.Now() + for { + select { + case <-done: + return + case <-ticker.C: + elapsedMs := int(time.Since(start) / time.Millisecond) + if elapsedMs >= estMs { + elapsedMs = estMs - 1 + } + // Intermediate ticks are best-effort: a panicking + // callback drops this tick but must not crash the + // process or leak the ticker goroutine. + func() { + defer func() { _ = recover() }() + progress(elapsedMs, estMs) + }() + } + } + }() + + var success bool + defer func() { + close(done) + <-stopped + if success { + progress(estMs, estMs) + } + }() + + if err := work(); err != nil { + return err + } + success = true + return nil +} + // ReadFlash reads data from flash memory. // Requires the stub loader to be running. // If progress is non-nil, it is called after each block is read with the diff --git a/pkg/espflasher/flasher_test.go b/pkg/espflasher/flasher_test.go index c49a55a..bf44dcd 100644 --- a/pkg/espflasher/flasher_test.go +++ b/pkg/espflasher/flasher_test.go @@ -490,7 +490,7 @@ func TestFlashMD5RequiresStub(t *testing.T) { mock := &mockConnection{} mock.stubMode = false // ROM mode f := &Flasher{conn: mock, chip: chipDefs[ChipESP32]} - _, err := f.GetFlashMD5(0, 1024) + _, err := f.GetFlashMD5(0, 1024, nil) if err == nil { t.Fatal("expected error when stub is not running") } diff --git a/pkg/espflasher/md5_test.go b/pkg/espflasher/md5_test.go new file mode 100644 index 0000000..ca69ad1 --- /dev/null +++ b/pkg/espflasher/md5_test.go @@ -0,0 +1,259 @@ +package espflasher + +import ( + "errors" + "runtime" + "testing" + "time" +) + +func TestTickMD5Progress(t *testing.T) { + interval := 2 * time.Millisecond + est := 50 * time.Millisecond + + var calls [][2]int + progress := func(current, total int) { + calls = append(calls, [2]int{current, total}) + } + work := func() error { + time.Sleep(6 * interval) + return nil + } + + if err := tickMD5(est, interval, progress, work); err != nil { + t.Fatalf("tickMD5 returned error: %v", err) + } + + if len(calls) < 2 { + t.Fatalf("expected at least one intermediate tick plus the final call, got %d calls", len(calls)) + } + + estMs := int(est / time.Millisecond) + last := calls[len(calls)-1] + if last[0] != estMs || last[1] != estMs { + t.Errorf("final call = %v, want (%d, %d)", last, estMs, estMs) + } + + prev := -1 + for i, c := range calls[:len(calls)-1] { + if c[1] != estMs { + t.Errorf("call %d total = %d, want %d", i, c[1], estMs) + } + if c[0] >= estMs { + t.Errorf("intermediate call %d current = %d, must be < total %d", i, c[0], estMs) + } + if c[0] < prev { + t.Errorf("call %d current = %d is not monotonically non-decreasing (prev %d)", i, c[0], prev) + } + prev = c[0] + } +} + +func TestTickMD5ErrorOmitsFinalCall(t *testing.T) { + interval := 2 * time.Millisecond + est := 50 * time.Millisecond + wantErr := errors.New("md5 failed") + + var calls [][2]int + progress := func(current, total int) { + calls = append(calls, [2]int{current, total}) + } + work := func() error { + time.Sleep(4 * interval) + return wantErr + } + + err := tickMD5(est, interval, progress, work) + if !errors.Is(err, wantErr) { + t.Fatalf("tickMD5 error = %v, want %v", err, wantErr) + } + + estMs := int(est / time.Millisecond) + for i, c := range calls { + if c[0] == estMs && c[1] == estMs { + t.Errorf("call %d reported completion (%d, %d) despite work returning an error", i, c[0], c[1]) + } + } +} + +func TestTickMD5JoinsGoroutine(t *testing.T) { + before := runtime.NumGoroutine() + + interval := time.Millisecond + est := 20 * time.Millisecond + work := func() error { + time.Sleep(5 * interval) + return nil + } + + if err := tickMD5(est, interval, func(int, int) {}, work); err != nil { + t.Fatalf("tickMD5 returned error: %v", err) + } + + // The ticker goroutine is joined synchronously inside tickMD5, so the + // count should already be back to baseline. Poll briefly anyway to absorb + // unrelated runtime bookkeeping goroutines. + deadline := time.Now().Add(200 * time.Millisecond) + for runtime.NumGoroutine() > before && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if got := runtime.NumGoroutine(); got > before { + t.Errorf("goroutine count after tickMD5 = %d, want <= %d (possible leak)", got, before) + } +} + +func TestTickMD5PanicJoinsGoroutineAndPropagates(t *testing.T) { + before := runtime.NumGoroutine() + + interval := time.Millisecond + est := 20 * time.Millisecond + wantPanic := "md5 blew up" + work := func() error { + time.Sleep(5 * interval) + panic(wantPanic) + } + + var calls [][2]int + progress := func(current, total int) { + calls = append(calls, [2]int{current, total}) + } + + func() { + defer func() { + r := recover() + if r == nil { + t.Fatal("expected tickMD5 to panic, but it did not") + } + if r != wantPanic { + t.Errorf("recovered panic = %v, want %v", r, wantPanic) + } + }() + _ = tickMD5(est, interval, progress, work) + }() + + estMs := int(est / time.Millisecond) + for i, c := range calls { + if c[0] == estMs && c[1] == estMs { + t.Errorf("call %d reported completion (%d, %d) despite work panicking", i, c[0], c[1]) + } + } + + // The ticker goroutine must be joined during panic unwinding, so the + // count should already be back to baseline. Poll briefly anyway to absorb + // unrelated runtime bookkeeping goroutines. + deadline := time.Now().Add(200 * time.Millisecond) + for runtime.NumGoroutine() > before && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if got := runtime.NumGoroutine(); got > before { + t.Errorf("goroutine count after tickMD5 panic = %d, want <= %d (possible leak)", got, before) + } +} + +func TestTickMD5IntermediatePanicDropsTickAndContinues(t *testing.T) { + before := runtime.NumGoroutine() + + interval := time.Millisecond + est := 50 * time.Millisecond + + // Panic only on an intermediate tick (not the first, not the final). + // The first intermediate call is allowed through, the second panics, + // and the final completion call must still be delivered normally. + var calls [][2]int + intermediate := 0 + estMs := int(est / time.Millisecond) + progress := func(current, total int) { + if current != estMs { // intermediate tick + intermediate++ + if intermediate == 2 { + panic("progress callback blew up on an intermediate tick") + } + } + calls = append(calls, [2]int{current, total}) + } + work := func() error { + time.Sleep(10 * interval) + return nil + } + + // A panic on an intermediate tick must not crash the process; tickMD5 + // must still return normally. + if err := tickMD5(est, interval, progress, work); err != nil { + t.Fatalf("tickMD5 returned error: %v", err) + } + + if len(calls) == 0 { + t.Fatal("expected at least the final progress call") + } + last := calls[len(calls)-1] + if last[0] != estMs || last[1] != estMs { + t.Errorf("final call = %v, want (%d, %d) — final tick must still fire after a dropped intermediate tick", last, estMs, estMs) + } + + // The ticker goroutine must still be joined after a dropped tick. + deadline := time.Now().Add(200 * time.Millisecond) + for runtime.NumGoroutine() > before && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if got := runtime.NumGoroutine(); got > before { + t.Errorf("goroutine count after intermediate-panic tickMD5 = %d, want <= %d (possible leak)", got, before) + } +} + +func TestGetFlashMD5NilProgressUnchanged(t *testing.T) { + before := runtime.NumGoroutine() + + mc := &mockConnection{ + stubMode: true, + flashMD5Func: func(addr, size uint32) ([]byte, error) { + return []byte{0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, + 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 0x10}, nil + }, + } + f := &Flasher{conn: mc, opts: DefaultOptions(), chip: chipDefs[ChipESP8266]} + + got, err := f.GetFlashMD5(0, 1024, nil) + if err != nil { + t.Fatalf("GetFlashMD5 returned error: %v", err) + } + want := "0102030405060708090a0b0c0d0e0f10" + if got != want { + t.Errorf("GetFlashMD5 = %q, want %q", got, want) + } + + deadline := time.Now().Add(200 * time.Millisecond) + for runtime.NumGoroutine() > before && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if gotN := runtime.NumGoroutine(); gotN > before { + t.Errorf("goroutine count after nil-progress GetFlashMD5 = %d, want <= %d (no ticker should run)", gotN, before) + } +} + +func TestGetFlashMD5ProgressWiring(t *testing.T) { + var calls [][2]int + mc := &mockConnection{ + stubMode: true, + flashMD5Func: func(addr, size uint32) ([]byte, error) { + return make([]byte, 16), nil + }, + } + f := &Flasher{conn: mc, opts: DefaultOptions(), chip: chipDefs[ChipESP8266]} + + size := uint32(512 * 1024) + _, err := f.GetFlashMD5(0, size, func(current, total int) { + calls = append(calls, [2]int{current, total}) + }) + if err != nil { + t.Fatalf("GetFlashMD5 returned error: %v", err) + } + if len(calls) == 0 { + t.Fatal("expected at least the final progress call") + } + + wantMs := int(md5TimeoutForSize(size) / time.Millisecond) + last := calls[len(calls)-1] + if last[0] != wantMs || last[1] != wantMs { + t.Errorf("final progress = %v, want (%d, %d)", last, wantMs, wantMs) + } +} diff --git a/pkg/espflasher/protocol.go b/pkg/espflasher/protocol.go index c8c29c4..d3cf58c 100644 --- a/pkg/espflasher/protocol.go +++ b/pkg/espflasher/protocol.go @@ -542,10 +542,7 @@ func (c *conn) flashMD5(addr, size uint32) ([]byte, error) { binary.LittleEndian.PutUint32(data[12:16], 0) // MD5 can take a while for large regions - timeout := md5Timeout - if size > 1024*1024 { - timeout = time.Duration(float64(md5Timeout) * float64(size) / float64(1024*1024)) - } + timeout := md5TimeoutForSize(size) result, err := c.checkCommand("flash MD5", cmdSPIFlashMD5, data, 0, timeout, 16) if err != nil { @@ -668,6 +665,15 @@ func eraseTimeoutForSize(size uint32) time.Duration { return t } +// md5TimeoutForSize calculates an appropriate timeout for MD5 hash +// operations, since larger regions can take a while to hash. +func md5TimeoutForSize(size uint32) time.Duration { + if size <= 1024*1024 { + return md5Timeout + } + return time.Duration(float64(md5Timeout) * float64(size) / float64(1024*1024)) +} + // flashWriteTimeoutForSize calculates an appropriate ack timeout for // compressed flash write/finish commands, scaled by data size using the // same per-MB rate as eraseTimeoutForSize. Unlike erase (which floors at