address comments

Author: elainezhao1

toolkit/tools/pkg/imagecustomizerlib/cosicommon.go

@@ -113,6 +113,7 @@ func buildCosiFile(sourceDir string, outputFile string, partitions []outputParti
 
 				veritySourcePath := path.Join(sourceDir, hashPartition.PartitionFilename)
 				imageDataEntry.VeritySource = veritySourcePath
+				break
 			}
 		}
 
@@ -177,41 +178,40 @@ func buildCosiFile(sourceDir string, outputFile string, partitions []outputParti
 }
 
 func addToCosi(data ImageBuildData, tw *tar.Writer) error {
-	imageFile, err := os.Open(data.Source)
-	if err != nil {
-		return fmt.Errorf("failed to open image file:\n%w", err)
-	}
-	defer imageFile.Close()
-
-	addFileToCosi(tw, data.Metadata.Image.Path, int64(data.Metadata.Image.CompressedSize), imageFile)
+	addFileToCosi(tw, data.Source, data.Metadata.Image)
 
 	if data.VeritySource != "" && data.Metadata.Verity != nil {
-		verityFile, err := os.Open(data.VeritySource)
+		err := addFileToCosi(tw, data.VeritySource, data.Metadata.Verity.Image)
 		if err != nil {
-			return fmt.Errorf("failed to open verity file:\n%w", err)
+			return fmt.Errorf("failed to add verity file to COSI:\n%w", err)
 		}
-		defer verityFile.Close()
-		addFileToCosi(tw, data.Metadata.Verity.Image.Path, int64(data.Metadata.Verity.Image.CompressedSize), verityFile)
 	}
 
 	return nil
 }
 
-func addFileToCosi(tw *tar.Writer, name string, size int64, file *os.File) error {
-	err := tw.WriteHeader(&tar.Header{
+func addFileToCosi(tw *tar.Writer, source string, image ImageFile) error {
+	file, err := os.Open(source)
+	if err != nil {
+		return fmt.Errorf("failed to open file :\n%w", err)
+	}
+	defer file.Close()
+
+	err = tw.WriteHeader(&tar.Header{
 		Typeflag: tar.TypeReg,
-		Name:     name,
-		Size:     size,
+		Name:     image.Path,
+		Size:     int64(image.CompressedSize),
 		Mode:     0o400,
 		Format:   tar.FormatPAX,
 	})
+
 	if err != nil {
-		return fmt.Errorf("failed to write tar header for file '%s':\n%w", name, err)
+		return fmt.Errorf("failed to write tar header for file '%s':\n%w", image.Path, err)
 	}
 
 	_, err = io.Copy(tw, file)
 	if err != nil {
-		return fmt.Errorf("failed to write image '%s' to COSI:\n%w", name, err)
+		return fmt.Errorf("failed to write image '%s' to COSI:\n%w", image.Path, err)
 	}
 
 	return nil