From 22336cde61ce14ade8a1fd814f2d7867b9455fe3 Mon Sep 17 00:00:00 2001
From: Chris Goller <goller@gmail.com>
Date: Mon, 30 Oct 2023 08:55:22 -0500
Subject: [PATCH] fix(buildx): use TempDir for configuration atomic writes

Previously, buildx would fail if extra files are in the
`.docker/instances` directory.  This could happen if
the program were to stop before the instance configuration
file was completely written leaving behind a temporary file.

This rewrites dockers's ioutils.AtomicWriteFile to create
the temporary file in the TempDir avoiding the problem.

Signed-off-by: Chris Goller <goller@gmail.com>
---
 pkg/cmd/docker/docker.go | 73 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 71 insertions(+), 2 deletions(-)

diff --git a/pkg/cmd/docker/docker.go b/pkg/cmd/docker/docker.go
index affbeb91..e6b6db5d 100644
--- a/pkg/cmd/docker/docker.go
+++ b/pkg/cmd/docker/docker.go
@@ -2,10 +2,12 @@ package docker
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"io"
 	"os"
 	"path"
+	"path/filepath"
 	"runtime"
 	"strings"
 
@@ -181,7 +183,11 @@ func runConfigureBuildx(ctx context.Context, dockerCli command.Cli, project, tok
 		return errors.Errorf("unknown project ID (run `depot init` or use --project or $DEPOT_PROJECT_ID)")
 	}
 
-	txn, release, err := storeutil.GetStore(dockerCli)
+	configStore, err := store.New(confutil.ConfigDir(dockerCli))
+	if err != nil {
+		return fmt.Errorf("unable to create docker configuration store: %w", err)
+	}
+	txn, release, err := configStore.Txn()
 	if err != nil {
 		return fmt.Errorf("unable to get docker store: %w", err)
 	}
@@ -263,7 +269,10 @@ func runConfigureBuildx(ctx context.Context, dockerCli command.Cli, project, tok
 		ng.Nodes[0], ng.Nodes[1] = ng.Nodes[1], ng.Nodes[0]
 	}
 
-	if err := txn.Save(ng); err != nil {
+	// DEPOT: we override the buildx Txn.Save() as its atomic write file
+	// can leave temporary files within the instance directory thus causing
+	// buildx to fail.
+	if err := DepotSaveNodes(confutil.ConfigDir(dockerCli), ng); err != nil {
 		return fmt.Errorf("unable to save node group: %w", err)
 	}
 
@@ -514,3 +523,63 @@ func CreateContainer(ctx context.Context, dockerCli command.Cli, projectName str
 
 	return nil
 }
+
+// DEPOT: we override the buildx Txn.Save() as its atomic write file
+// can leave temporary files within the instance directory thus causing
+// buildx to fail.
+func DepotSaveNodes(configDir string, ng *store.NodeGroup) (err error) {
+	name, err := store.ValidateName(ng.Name)
+	if err != nil {
+		return err
+	}
+
+	octets, err := json.Marshal(ng)
+	if err != nil {
+		return err
+	}
+
+	instancesDir := filepath.Join(configDir, "instances")
+	fileName := filepath.Join(instancesDir, name)
+
+	// DEPOT: this is the key change for saving the nodes.
+	// Previously, it would save in the instances directory, but
+	// those files would then be read by the Txn.List()/Txn.NodeGroupByName()
+	// methods and thus would fail.
+	//
+	// Instead, we save the file to the TempDir and then rename it.
+	// CreateTemp creates a file with 0600 perms.
+	f, err := os.CreateTemp("", ".tmp-"+filepath.Base(fileName))
+	if err != nil {
+		return err
+	}
+
+	// Require that the file be removed on error.
+	defer func() {
+		if err != nil {
+			_ = f.Close()
+			_ = os.Remove(f.Name())
+		}
+	}()
+
+	n, err := f.Write(octets)
+	if err != nil {
+		return err
+	}
+
+	if n != len(octets) {
+		err = io.ErrShortWrite
+		return err
+	}
+
+	err = f.Sync()
+	if err != nil {
+		return err
+	}
+
+	err = f.Close()
+	if err != nil {
+		return err
+	}
+
+	return os.Rename(f.Name(), fileName)
+}