diff --git a/reflash/server.go b/reflash/server.go index cbb286a..bba9b11 100644 --- a/reflash/server.go +++ b/reflash/server.go @@ -578,7 +578,11 @@ func uploadStart(w http.ResponseWriter, r *http.Request) { timeStart = time.Now() logInfo("Starting upload at " + timeStart.Format("15:04:05")) logInfo("Filename: " + state.Filename) - os.Create(images_folder + "/" + state.Filename) + f, err := os.Create(images_folder + "/" + state.Filename) + if err != nil { + log.Fatal(err) + } + state.File = f sendResponse(w, nil) } @@ -674,25 +678,26 @@ func uploadChunk(w http.ResponseWriter, r *http.Request) { reqBody, _ := io.ReadAll(r.Body) json.Unmarshal(reqBody, &chunk) - decoded, err := base64.StdEncoding.DecodeString(chunk.Encoded[37:len(chunk.Encoded)]) - - path := images_folder + "/" + state.Filename - if state.State == CANCELLED { response := map[string]bool{"success": false} json.NewEncoder(w).Encode(response) return } - f, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0644) + decoded, err := base64.StdEncoding.DecodeString(chunk.Encoded[37:len(chunk.Encoded)]) if err != nil { - log.Fatal(err) - } - if _, err := f.Write(decoded); err != nil { - log.Fatal(err) + http.Error(w, "Failed to decode base64", http.StatusBadRequest) + return } - if err := f.Close(); err != nil { - log.Fatal(err) + + // state.File is opened once in uploadStart and closed in + // uploadFinish/uploadCancel - opening, writing and closing it on + // every chunk here was real per-chunk overhead (syscalls plus an + // implicit flush on close), especially costly against the + // FAT-formatted USB target. + if _, err := state.File.Write(decoded); err != nil { + http.Error(w, "Failed to write chunk to file", http.StatusInternalServerError) + return } state.BytesNow += len(decoded) state.Progress = float64(state.BytesNow) * 100 / float64(state.BytesTotal) @@ -702,6 +707,12 @@ func uploadChunk(w http.ResponseWriter, r *http.Request) { } func uploadFinish(w http.ResponseWriter, r *http.Request) { + if state.File != nil { + if err := state.File.Close(); err != nil { + log.Fatal(err) + } + state.File = nil + } mountUsb(MODE_RO) duration := time.Since(timeStart) logInfo(fmt.Sprintf("Upload finished in %d minutes and %d seconds", int(duration.Minutes()), int(duration.Seconds())%60)) diff --git a/reflash/server_test.go b/reflash/server_test.go index 7e6bc4d..c32f64f 100644 --- a/reflash/server_test.go +++ b/reflash/server_test.go @@ -1,6 +1,8 @@ package main import ( + "bytes" + "encoding/base64" "encoding/json" "net/http/httptest" "os" @@ -393,3 +395,65 @@ func TestGetUncompressedSizeRealXz(t *testing.T) { t.Errorf("getUncompressedSize = %d, want %d", got, size) } } + +// End-to-end against the real handlers: drives upload_start -> several +// upload_chunk calls -> upload_finish exactly like the client does, and +// checks the file on disk matches byte-for-byte. Exercises the real +// state.File handle kept open across chunks (rather than the old +// open/write/close-per-chunk code), so this would catch truncation, +// overwriting, or interleaving bugs a build-only check can't. +func TestUploadChunkRoundTrip(t *testing.T) { + setupTest(t) + state = &State{State: IDLE} + + filename := "roundtrip-test.img.xz" + payload := bytes.Repeat([]byte("The quick brown fox jumps over the lazy dog. "), 5000) // ~225KB + + startBody, _ := json.Marshal(map[string]any{ + "filename": filename, + "size": len(payload), + "start_time": 0, + }) + w := httptest.NewRecorder() + uploadStart(w, httptest.NewRequest("PUT", "/api/upload_start", bytes.NewReader(startBody))) + if state.File == nil { + t.Fatal("uploadStart did not open state.File") + } + + const chunkSize = 10 * 1024 // small relative to payload so this genuinely exercises multiple sequential chunk writes, not just one + for i := 0; i < len(payload); i += chunkSize { + end := i + chunkSize + if end > len(payload) { + end = len(payload) + } + encoded := "data:application/octet-stream;base64," + base64.StdEncoding.EncodeToString(payload[i:end]) + chunkBody, _ := json.Marshal(map[string]string{"chunk": encoded}) + w := httptest.NewRecorder() + uploadChunk(w, httptest.NewRequest("POST", "/api/upload_chunk", bytes.NewReader(chunkBody))) + + var resp map[string]bool + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("chunk at offset %d: bad response body: %v", i, err) + } + if !resp["success"] { + t.Fatalf("chunk at offset %d: server reported failure", i) + } + } + + uploadFinish(httptest.NewRecorder(), httptest.NewRequest("PUT", "/api/upload_finish", nil)) + + if state.File != nil { + t.Error("uploadFinish did not clear state.File") + } + + got, err := os.ReadFile(filepath.Join(images_folder, filename)) + if err != nil { + t.Fatalf("reading uploaded file: %v", err) + } + if !bytes.Equal(got, payload) { + t.Errorf("uploaded file mismatch: got %d bytes, want %d bytes", len(got), len(payload)) + } + if state.BytesNow != len(payload) { + t.Errorf("state.BytesNow = %d, want %d", state.BytesNow, len(payload)) + } +}