fix: bandwidth saved shows 0B

aloks98 · aloks98 · commit 46fc260b5e01 · 2025-12-16T01:48:04.000+05:30
diff --git a/backend/internal/db/sqlite.go b/backend/internal/db/sqlite.go
@@ -429,3 +429,24 @@ func (db *DB) GetISOByComposite(name, version, arch, edition, fileType string) (
 	}
 	return iso, nil
 }
+
+// ListISOsWithMissingSize returns ISOs that are complete but have size_bytes = 0.
+func (db *DB) ListISOsWithMissingSize() ([]models.ISO, error) {
+	query := fmt.Sprintf("SELECT %s FROM isos WHERE status = 'complete' AND size_bytes = 0", isoSelectFields)
+	rows, err := db.conn.Query(query)
+	if err != nil {
+		return nil, fmt.Errorf("failed to list ISOs with missing size: %w", err)
+	}
+	defer rows.Close()
+
+	var isos []models.ISO
+	for rows.Next() {
+		iso, err := scanISO(rows)
+		if err != nil {
+			return nil, fmt.Errorf("failed to scan ISO row: %w", err)
+		}
+		isos = append(isos, *iso)
+	}
+
+	return isos, rows.Err()
+}
diff --git a/backend/internal/db/sqlite_test.go b/backend/internal/db/sqlite_test.go
@@ -762,3 +762,112 @@ func TestConcurrentOperations(t *testing.T) {
 		t.Error("ISO corrupted after concurrent operations")
 	}
 }
+
+func TestListISOsWithMissingSize(t *testing.T) {
+	db, cleanup := setupTestDB(t)
+	defer cleanup()
+
+	// Create ISOs with various states
+	// ISO 1: complete with size_bytes = 0 (should be returned)
+	iso1 := createTestISO()
+	iso1.ID = "iso-missing-size-1"
+	iso1.Name = "missing-size-1"
+	iso1.Status = models.StatusComplete
+	iso1.SizeBytes = 0
+	if err := db.CreateISO(iso1); err != nil {
+		t.Fatalf("Failed to create iso1: %v", err)
+	}
+
+	// ISO 2: complete with size_bytes > 0 (should NOT be returned)
+	iso2 := createTestISO()
+	iso2.ID = "iso-has-size"
+	iso2.Name = "has-size"
+	iso2.Status = models.StatusComplete
+	iso2.SizeBytes = 1000000
+	if err := db.CreateISO(iso2); err != nil {
+		t.Fatalf("Failed to create iso2: %v", err)
+	}
+
+	// ISO 3: pending with size_bytes = 0 (should NOT be returned - not complete)
+	iso3 := createTestISO()
+	iso3.ID = "iso-pending"
+	iso3.Name = "pending"
+	iso3.Status = models.StatusPending
+	iso3.SizeBytes = 0
+	if err := db.CreateISO(iso3); err != nil {
+		t.Fatalf("Failed to create iso3: %v", err)
+	}
+
+	// ISO 4: failed with size_bytes = 0 (should NOT be returned - not complete)
+	iso4 := createTestISO()
+	iso4.ID = "iso-failed"
+	iso4.Name = "failed"
+	iso4.Status = models.StatusFailed
+	iso4.SizeBytes = 0
+	if err := db.CreateISO(iso4); err != nil {
+		t.Fatalf("Failed to create iso4: %v", err)
+	}
+
+	// ISO 5: complete with size_bytes = 0 (should be returned)
+	iso5 := createTestISO()
+	iso5.ID = "iso-missing-size-2"
+	iso5.Name = "missing-size-2"
+	iso5.Status = models.StatusComplete
+	iso5.SizeBytes = 0
+	if err := db.CreateISO(iso5); err != nil {
+		t.Fatalf("Failed to create iso5: %v", err)
+	}
+
+	// Test ListISOsWithMissingSize
+	isos, err := db.ListISOsWithMissingSize()
+	if err != nil {
+		t.Fatalf("ListISOsWithMissingSize() failed: %v", err)
+	}
+
+	// Should return exactly 2 ISOs (iso1 and iso5)
+	if len(isos) != 2 {
+		t.Errorf("Expected 2 ISOs with missing size, got %d", len(isos))
+	}
+
+	// Verify the returned ISOs are the correct ones
+	foundIDs := make(map[string]bool)
+	for _, iso := range isos {
+		foundIDs[iso.ID] = true
+		// All returned ISOs should be complete with size_bytes = 0
+		if iso.Status != models.StatusComplete {
+			t.Errorf("Returned ISO %s has status %s, expected complete", iso.ID, iso.Status)
+		}
+		if iso.SizeBytes != 0 {
+			t.Errorf("Returned ISO %s has size_bytes %d, expected 0", iso.ID, iso.SizeBytes)
+		}
+	}
+
+	if !foundIDs["iso-missing-size-1"] {
+		t.Error("Expected iso-missing-size-1 to be returned")
+	}
+	if !foundIDs["iso-missing-size-2"] {
+		t.Error("Expected iso-missing-size-2 to be returned")
+	}
+}
+
+func TestListISOsWithMissingSize_Empty(t *testing.T) {
+	db, cleanup := setupTestDB(t)
+	defer cleanup()
+
+	// Create only ISOs that should NOT be returned
+	iso := createTestISO()
+	iso.Status = models.StatusComplete
+	iso.SizeBytes = 500000
+	if err := db.CreateISO(iso); err != nil {
+		t.Fatalf("Failed to create ISO: %v", err)
+	}
+
+	isos, err := db.ListISOsWithMissingSize()
+	if err != nil {
+		t.Fatalf("ListISOsWithMissingSize() failed: %v", err)
+	}
+
+	if len(isos) != 0 {
+		t.Errorf("Expected 0 ISOs with missing size, got %d", len(isos))
+	}
+}
diff --git a/backend/internal/download/worker.go b/backend/internal/download/worker.go
@@ -90,6 +90,16 @@ func (w *Worker) Process(ctx context.Context, iso *models.ISO) error {
 		return fmt.Errorf("failed to move file to final location: %w", err)
 	}
 
+	// Update size_bytes from actual file size if not set (e.g., server didn't send Content-Length)
+	if iso.SizeBytes == 0 {
+		if fi, err := os.Stat(finalFile); err == nil {
+			iso.SizeBytes = fi.Size()
+			if err := w.db.UpdateISOSize(iso.ID, iso.SizeBytes); err != nil {
+				slog.Warn("failed to update ISO size from file", slog.Any("error", err))
+			}
+		}
+	}
+
 	// Download and save checksum file alongside ISO (after file is moved)
 	if iso.ChecksumURL != "" {
 		checksumFile := pathutil.ConstructChecksumPath(finalFile, iso.ChecksumType)
diff --git a/backend/internal/download/worker_test.go b/backend/internal/download/worker_test.go
@@ -516,3 +516,85 @@ func TestWorkerTempFileCleanup(t *testing.T) {
 		t.Errorf("Temp file should be cleaned up: %s", tmpFile)
 	}
 }
+
+// TestWorkerDownloadNoContentLength tests that size_bytes is set from actual file
+// when server doesn't send Content-Length header.
+func TestWorkerDownloadNoContentLength(t *testing.T) {
+	worker, database, isoDir, cleanup := setupTestWorker(t)
+	defer cleanup()
+
+	// Create test HTTP server WITHOUT Content-Length header
+	testContent := []byte("test file content without content-length header")
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Explicitly remove Content-Length by using chunked transfer
+		w.Header().Set("Transfer-Encoding", "chunked")
+		w.WriteHeader(http.StatusOK)
+		w.Write(testContent)
+	}))
+	defer server.Close()
+
+	// Create test ISO
+	iso := &models.ISO{
+		ID:          uuid.New().String(),
+		Name:        "test-no-cl",
+		Version:     "1.0",
+		Arch:        "x86_64",
+		FileType:    "iso",
+		DownloadURL: server.URL,
+		Status:      models.StatusPending,
+		Progress:    0,
+		SizeBytes:   0, // Explicitly set to 0
+		CreatedAt:   time.Now(),
+	}
+	iso.ComputeFields()
+	database.CreateISO(iso)
+
+	// Verify initial size is 0
+	initialISO, _ := database.GetISO(iso.ID)
+	if initialISO.SizeBytes != 0 {
+		t.Errorf("Initial SizeBytes should be 0, got: %d", initialISO.SizeBytes)
+	}
+
+	// Process download
+	ctx := context.Background()
+	err := worker.Process(ctx, iso)
+	if err != nil {
+		t.Fatalf("Process failed: %v", err)
+	}
+
+	// Verify file was downloaded
+	finalPath := filepath.Join(isoDir, iso.FilePath)
+	fi, err := os.Stat(finalPath)
+	if os.IsNotExist(err) {
+		t.Fatalf("Downloaded file does not exist at %s", finalPath)
+	}
+
+	// Verify file content matches
+	content, err := os.ReadFile(finalPath)
+	if err != nil {
+		t.Fatalf("Failed to read downloaded file: %v", err)
+	}
+	if !bytes.Equal(content, testContent) {
+		t.Errorf("File content mismatch: got %q, want %q", content, testContent)
+	}
+
+	// THIS IS THE KEY TEST: Verify size_bytes was backfilled from actual file size
+	updatedISO, err := database.GetISO(iso.ID)
+	if err != nil {
+		t.Fatalf("Failed to get updated ISO: %v", err)
+	}
+
+	if updatedISO.Status != models.StatusComplete {
+		t.Errorf("Status should be 'complete', got: %s", updatedISO.Status)
+	}
+
+	// size_bytes should match actual file size even without Content-Length
+	expectedSize := fi.Size()
+	if updatedISO.SizeBytes != expectedSize {
+		t.Errorf("SizeBytes should be %d (actual file size), got: %d", expectedSize, updatedISO.SizeBytes)
+	}
+
+	if updatedISO.SizeBytes != int64(len(testContent)) {
+		t.Errorf("SizeBytes should be %d (content length), got: %d", len(testContent), updatedISO.SizeBytes)
+	}
+}
diff --git a/backend/main.go b/backend/main.go
@@ -58,6 +58,9 @@ func main() {
 	defer database.Close()
 	log.Info("database initialized", slog.String("db_path", dbPath))
 
+	// Backfill missing ISO sizes from actual files on disk
+	backfillISOSizes(database, isoDir, log)
+
 	// Initialize WebSocket hub
 	wsHub := ws.NewHub()
 	go wsHub.Run()
@@ -135,3 +138,54 @@ func main() {
 
 	log.Info("server stopped successfully")
 }
+
+// backfillISOSizes updates size_bytes for complete ISOs that have size_bytes = 0
+// by reading the actual file size from disk. This handles ISOs that were downloaded
+// when the server didn't send a Content-Length header.
+func backfillISOSizes(database *db.DB, isoDir string, log *slog.Logger) {
+	isos, err := database.ListISOsWithMissingSize()
+	if err != nil {
+		log.Warn("failed to list ISOs with missing size", slog.Any("error", err))
+		return
+	}
+
+	if len(isos) == 0 {
+		return
+	}
+
+	log.Info("backfilling missing ISO sizes", slog.Int("count", len(isos)))
+
+	updated := 0
+	for _, iso := range isos {
+		filePath := pathutil.ConstructISOPath(isoDir, iso.FilePath)
+		fi, err := os.Stat(filePath)
+		if err != nil {
+			log.Warn("failed to stat ISO file for size backfill",
+				slog.String("iso_id", iso.ID),
+				slog.String("path", filePath),
+				slog.Any("error", err),
+			)
+			continue
+		}
+
+		size := fi.Size()
+		if err := database.UpdateISOSize(iso.ID, size); err != nil {
+			log.Warn("failed to update ISO size",
+				slog.String("iso_id", iso.ID),
+				slog.Any("error", err),
+			)
+			continue
+		}
+
+		updated++
+		log.Debug("backfilled ISO size",
+			slog.String("iso_id", iso.ID),
+			slog.String("name", iso.Name),
+			slog.Int64("size_bytes", size),
+		)
+	}
+
+	if updated > 0 {
+		log.Info("backfilled ISO sizes", slog.Int("updated", updated))
+	}
+}
