From 8d1ba540f99e896984b8c556b44f6b3d7a6d42e6 Mon Sep 17 00:00:00 2001
From: Adam Hughes <stickmanica@gmail.com>
Date: Mon, 19 Jul 2021 23:21:25 +0000
Subject: [PATCH] fix: return usable FileImage from CreateContainer

Do not close fimg.Fp in CreateContainer so that fimg is usable by
caller.
---
 internal/app/siftool/modif.go      |  8 ++++++--
 pkg/integrity/testdata/gen_sifs.go | 21 ++++++++++-----------
 pkg/sif/create.go                  |  9 ++++++++-
 pkg/sif/create_test.go             | 14 ++++++++++++--
 4 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/internal/app/siftool/modif.go b/internal/app/siftool/modif.go
index f6111ce4..53597742 100644
--- a/internal/app/siftool/modif.go
+++ b/internal/app/siftool/modif.go
@@ -29,8 +29,12 @@ func New(path string) error {
 		ID:         id,
 	}
 
-	_, err = sif.CreateContainer(cinfo)
-	return err
+	fimg, err := sif.CreateContainer(cinfo)
+	if err != nil {
+		return err
+	}
+
+	return fimg.UnloadContainer()
 }
 
 // AddOptions contains the options when adding a section to a SIF file.
diff --git a/pkg/integrity/testdata/gen_sifs.go b/pkg/integrity/testdata/gen_sifs.go
index e68ce982..1c3efc35 100755
--- a/pkg/integrity/testdata/gen_sifs.go
+++ b/pkg/integrity/testdata/gen_sifs.go
@@ -17,10 +17,10 @@ import (
 	"golang.org/x/crypto/openpgp"
 )
 
-func createImage(path string, dis []sif.DescriptorInput) error {
+func createImage(path string, dis []sif.DescriptorInput) (*sif.FileImage, error) {
 	id, err := uuid.NewV4()
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	ci := sif.CreateInfo{
@@ -31,8 +31,7 @@ func createImage(path string, dis []sif.DescriptorInput) error {
 		InputDescr: dis,
 	}
 
-	_, err = sif.CreateContainer(ci)
-	return err
+	return sif.CreateContainer(ci)
 }
 
 func getEntity() (*openpgp.Entity, error) {
@@ -146,18 +145,18 @@ func generateImages() error {
 	for _, image := range images {
 		path := filepath.Join("images", image.path)
 
-		if err := createImage(path, image.dis); err != nil {
-			return err
-		}
-
-		f, err := sif.LoadContainer(path, false)
+		f, err := createImage(path, image.dis)
 		if err != nil {
 			return err
 		}
-		defer f.UnloadContainer() // nolint:errcheck
+		defer func() {
+			if err := f.UnloadContainer(); err != nil {
+				log.Printf("failed to unload container: %v", err)
+			}
+		}()
 
 		if image.sign {
-			s, err := integrity.NewSigner(&f, image.signOpts...)
+			s, err := integrity.NewSigner(f, image.signOpts...)
 			if err != nil {
 				return err
 			}
diff --git a/pkg/sif/create.go b/pkg/sif/create.go
index 98ec9264..6aa55bd8 100644
--- a/pkg/sif/create.go
+++ b/pkg/sif/create.go
@@ -219,6 +219,9 @@ func writeHeader(fimg *FileImage) error {
 // CreateContainer is responsible for the creation of a new SIF container
 // file. It takes the creation information specification as input
 // and produces an output file as specified in the input data.
+//
+// On success, a FileImage is returned. The caller must call UnloadContainer
+// to ensure resources are released.
 func CreateContainer(cinfo CreateInfo) (fimg *FileImage, err error) {
 	fimg = &FileImage{}
 	fimg.DescrArr = make([]Descriptor, DescrNumEntries)
@@ -241,7 +244,11 @@ func CreateContainer(cinfo CreateInfo) (fimg *FileImage, err error) {
 	if err != nil {
 		return nil, fmt.Errorf("container file creation failed: %s", err)
 	}
-	defer fimg.Fp.Close()
+	defer func() {
+		if err != nil {
+			fimg.Fp.Close()
+		}
+	}()
 
 	// set file pointer to start of data section */
 	if _, err = fimg.Fp.Seek(DataStartOffset, 0); err != nil {
diff --git a/pkg/sif/create_test.go b/pkg/sif/create_test.go
index 99e2fadb..48e3ff2b 100644
--- a/pkg/sif/create_test.go
+++ b/pkg/sif/create_test.go
@@ -79,10 +79,15 @@ func TestCreateContainer(t *testing.T) {
 	}
 
 	// test container creation without any input descriptors
-	if _, err := CreateContainer(cinfo); err != nil {
+	fimg, err := CreateContainer(cinfo)
+	if err != nil {
 		t.Error("CreateContainer(cinfo): should allow empty input descriptor list")
 	}
 
+	if err := fimg.UnloadContainer(); err != nil {
+		t.Error(err)
+	}
+
 	// data we need to create a definition file descriptor
 	definput := DescriptorInput{
 		Datatype: DataDeffile,
@@ -141,9 +146,14 @@ func TestCreateContainer(t *testing.T) {
 	cinfo.InputDescr = append(cinfo.InputDescr, parinput)
 
 	// test container creation with two partition input descriptors
-	if _, err := CreateContainer(cinfo); err != nil {
+	fimg, err = CreateContainer(cinfo)
+	if err != nil {
 		t.Errorf("CreateContainer(cinfo): CreateContainer(): %s", err)
 	}
+
+	if err := fimg.UnloadContainer(); err != nil {
+		t.Error(err)
+	}
 }
 
 func TestAddDelObject(t *testing.T) {