From 306f1854534c457fa4dda3880ff7cc82f6a85460 Mon Sep 17 00:00:00 2001
From: Mikhail Mazurskiy <126021+ash2k@users.noreply.github.com>
Date: Fri, 5 Jul 2024 20:36:35 +1000
Subject: [PATCH] Fix issues in BuildImage() (#2626)

* Don't accumulate errors

The code already logs them so no need to return them all to the caller - the calling code will likely log them too. This will result in the same errors being logged twice. This is confusing for the user.

* RetryNotify -> RetryNotifyWithData

* Do not reuse buildOptions.Context

buildOptions.Context cannot be reused because it might have been partially consumed by the failed ImageBuild(). So, the next invocation will send a partial context, which will not work correctly.

* Always close the response body

* Always close the build context

This ensures all resources are released properly.
---
 container.go | 14 ++++++++------
 docker.go    | 38 ++++++++++++++++++++++----------------
 2 files changed, 30 insertions(+), 22 deletions(-)

diff --git a/container.go b/container.go
index c08d45624d..30d0ced2b0 100644
--- a/container.go
+++ b/container.go
@@ -209,6 +209,7 @@ func (c *ContainerRequest) Validate() error {
 }
 
 // GetContext retrieve the build context for the request
+// Must be closed when no longer needed.
 func (c *ContainerRequest) GetContext() (io.Reader, error) {
 	var includes []string = []string{"."}
 
@@ -356,12 +357,6 @@ func (c *ContainerRequest) BuildOptions() (types.ImageBuildOptions, error) {
 	buildOptions.BuildArgs = c.GetBuildArgs()
 	buildOptions.Dockerfile = c.GetDockerfile()
 
-	buildContext, err := c.GetContext()
-	if err != nil {
-		return buildOptions, err
-	}
-	buildOptions.Context = buildContext
-
 	// Make sure the auth configs from the Dockerfile are set right after the user-defined build options.
 	authsFromDockerfile := getAuthConfigsFromDockerfile(c)
 
@@ -400,6 +395,13 @@ func (c *ContainerRequest) BuildOptions() (types.ImageBuildOptions, error) {
 		buildOptions.Labels = core.DefaultLabels(core.SessionID())
 	}
 
+	// Do this as late as possible to ensure we don't leak the context on error/panic.
+	buildContext, err := c.GetContext()
+	if err != nil {
+		return buildOptions, err
+	}
+	buildOptions.Context = buildContext
+
 	return buildOptions, nil
 }
 
diff --git a/docker.go b/docker.go
index d47a2d2681..30e32c44a4 100644
--- a/docker.go
+++ b/docker.go
@@ -882,26 +882,26 @@ var _ ContainerProvider = (*DockerProvider)(nil)
 
 // BuildImage will build and image from context and Dockerfile, then return the tag
 func (p *DockerProvider) BuildImage(ctx context.Context, img ImageBuildInfo) (string, error) {
-	buildOptions, err := img.BuildOptions()
-	if err != nil {
-		return "", err
-	}
+	var buildOptions types.ImageBuildOptions
+	resp, err := backoff.RetryNotifyWithData(
+		func() (types.ImageBuildResponse, error) {
+			var err error
+			buildOptions, err = img.BuildOptions()
+			if err != nil {
+				return types.ImageBuildResponse{}, backoff.Permanent(err)
+			}
+			defer tryClose(buildOptions.Context) // release resources in any case
 
-	var buildError error
-	var resp types.ImageBuildResponse
-	err = backoff.RetryNotify(
-		func() error {
-			resp, err = p.client.ImageBuild(ctx, buildOptions.Context, buildOptions)
+			resp, err := p.client.ImageBuild(ctx, buildOptions.Context, buildOptions)
 			if err != nil {
-				buildError = errors.Join(buildError, err)
 				if isPermanentClientError(err) {
-					return backoff.Permanent(err)
+					return types.ImageBuildResponse{}, backoff.Permanent(err)
 				}
-				return err
+				return types.ImageBuildResponse{}, err
 			}
 			defer p.Close()
 
-			return nil
+			return resp, nil
 		},
 		backoff.WithContext(backoff.NewExponentialBackOff(), ctx),
 		func(err error, duration time.Duration) {
@@ -909,8 +909,9 @@ func (p *DockerProvider) BuildImage(ctx context.Context, img ImageBuildInfo) (st
 		},
 	)
 	if err != nil {
-		return "", errors.Join(buildError, err)
+		return "", err
 	}
+	defer resp.Body.Close()
 
 	if img.ShouldPrintBuildLog() {
 		termFd, isTerm := term.GetFdInfo(os.Stderr)
@@ -927,8 +928,6 @@ func (p *DockerProvider) BuildImage(ctx context.Context, img ImageBuildInfo) (st
 		return "", err
 	}
 
-	_ = resp.Body.Close()
-
 	// the first tag is the one we want
 	return buildOptions.Tags[0], nil
 }
@@ -1664,3 +1663,10 @@ func isPermanentClientError(err error) bool {
 	}
 	return false
 }
+
+func tryClose(r io.Reader) {
+	rc, ok := r.(io.Closer)
+	if ok {
+		_ = rc.Close()
+	}
+}