From 972c3797550a1013706dc78aab4a1edd84752bfc Mon Sep 17 00:00:00 2001
From: Jacob Gillespie <jacobwgillespie@gmail.com>
Date: Wed, 14 Aug 2024 18:07:18 +0100
Subject: [PATCH] Cleanup temporary certificate files after build

---
 cmd/buildctl/main.go    |  3 +++
 cmd/depot/main.go       |  3 +++
 pkg/cleanup/tmpfiles.go | 15 +++++++++++++++
 pkg/machine/machine.go  |  4 ++++
 4 files changed, 25 insertions(+)
 create mode 100644 pkg/cleanup/tmpfiles.go

diff --git a/cmd/buildctl/main.go b/cmd/buildctl/main.go
index be7ca326..e0977b66 100644
--- a/cmd/buildctl/main.go
+++ b/cmd/buildctl/main.go
@@ -5,6 +5,7 @@ import (
 	"os"
 
 	"github.com/depot/cli/internal/build"
+	"github.com/depot/cli/pkg/cleanup"
 	"github.com/depot/cli/pkg/cmd/buildctl"
 	"github.com/getsentry/sentry-go"
 )
@@ -26,6 +27,8 @@ func runMain() int {
 		}
 	}
 
+	defer cleanup.CleanupTmpfiles()
+
 	err := buildctl.NewBuildctl().Execute()
 	if err != nil {
 		return 1
diff --git a/cmd/depot/main.go b/cmd/depot/main.go
index 58614ba8..189f83d1 100644
--- a/cmd/depot/main.go
+++ b/cmd/depot/main.go
@@ -12,6 +12,7 @@ import (
 	"github.com/depot/cli/internal/build"
 	"github.com/depot/cli/internal/update"
 	"github.com/depot/cli/pkg/api"
+	"github.com/depot/cli/pkg/cleanup"
 	"github.com/depot/cli/pkg/cmd/root"
 	"github.com/depot/cli/pkg/config"
 	"github.com/depot/cli/pkg/helpers"
@@ -82,6 +83,8 @@ func runMain() int {
 		}
 	}
 
+	defer cleanup.CleanupTmpfiles()
+
 	buildVersion := build.Version
 	buildDate := build.Date
 
diff --git a/pkg/cleanup/tmpfiles.go b/pkg/cleanup/tmpfiles.go
new file mode 100644
index 00000000..eaab259b
--- /dev/null
+++ b/pkg/cleanup/tmpfiles.go
@@ -0,0 +1,15 @@
+package cleanup
+
+import "os"
+
+var tmpfiles = []string{}
+
+func RegisterTmpfile(filename string) {
+	tmpfiles = append(tmpfiles, filename)
+}
+
+func CleanupTmpfiles() {
+	for _, filename := range tmpfiles {
+		_ = os.Remove(filename)
+	}
+}
diff --git a/pkg/machine/machine.go b/pkg/machine/machine.go
index 261dd4ac..c4602343 100644
--- a/pkg/machine/machine.go
+++ b/pkg/machine/machine.go
@@ -11,6 +11,7 @@ import (
 
 	"connectrpc.com/connect"
 	"github.com/depot/cli/pkg/api"
+	"github.com/depot/cli/pkg/cleanup"
 	cliv1 "github.com/depot/cli/pkg/proto/depot/cli/v1"
 	"github.com/depot/cli/pkg/proto/depot/cli/v1/cliv1connect"
 	"github.com/moby/buildkit/client"
@@ -180,6 +181,7 @@ func (m *Machine) Client(ctx context.Context) (*client.Client, error) {
 			return nil, errors.Wrap(err, "failed to write cert to temp file")
 		}
 		cert := file.Name()
+		cleanup.RegisterTmpfile(cert)
 
 		file, err = os.CreateTemp("", "depot-key")
 		if err != nil {
@@ -191,6 +193,7 @@ func (m *Machine) Client(ctx context.Context) (*client.Client, error) {
 			return nil, errors.Wrap(err, "failed to write key to temp file")
 		}
 		key := file.Name()
+		cleanup.RegisterTmpfile(key)
 
 		file, err = os.CreateTemp("", "depot-ca-cert")
 		if err != nil {
@@ -202,6 +205,7 @@ func (m *Machine) Client(ctx context.Context) (*client.Client, error) {
 			return nil, errors.Wrap(err, "failed to write CA cert to temp file")
 		}
 		caCert := file.Name()
+		cleanup.RegisterTmpfile(caCert)
 
 		opts = append(opts, client.WithCredentials(m.ServerName, caCert, cert, key))
 	}