fix: Extract the plugin binary atomically

Author: erezrokah

managedplugin/download.go

@@ -190,6 +190,8 @@ func doDownloadPluginFromHub(ctx context.Context, logger zerolog.Logger, c *clou
 		return errors.New("failed to get plugin metadata from hub: empty location from response")
 	}
 	pluginZipPath := ops.LocalPath + ".zip"
+	defer os.Remove(pluginZipPath)
+
 	writtenChecksum, err := downloadFile(ctx, pluginZipPath, location, dops)
 	if err != nil {
 		return fmt.Errorf("failed to download plugin: %w", err)
@@ -201,29 +203,46 @@ func doDownloadPluginFromHub(ctx context.Context, logger zerolog.Logger, c *clou
 		return fmt.Errorf("checksum mismatch: expected %s, got %s", pluginAsset.Checksum, writtenChecksum)
 	}
 
-	archive, err := zip.OpenReader(pluginZipPath)
+	pathInArchive := fmt.Sprintf("plugin-%s-%s-%s-%s", ops.PluginName, ops.PluginVersion, runtime.GOOS, runtime.GOARCH)
+	return extractPluginBinary(pluginZipPath, pathInArchive, ops.LocalPath)
+}
+
+// extractPluginBinary writes the binary to a temporary file and renames it into
+// place, so a failure part way through never leaves a truncated binary that the
+// next run treats as a cached plugin.
+func extractPluginBinary(archivePath, pathInArchive, localPath string) error {
+	archive, err := zip.OpenReader(archivePath)
 	if err != nil {
 		return fmt.Errorf("failed to open plugin archive: %w", err)
 	}
 	defer archive.Close()
 
-	fileInArchive, err := archive.Open(fmt.Sprintf("plugin-%s-%s-%s-%s", ops.PluginName, ops.PluginVersion, runtime.GOOS, runtime.GOARCH))
+	fileInArchive, err := archive.Open(pathInArchive)
 	if err != nil {
-		return fmt.Errorf("failed to open plugin archive: %w", err)
+		return fmt.Errorf("failed to open plugin archive %s: %w", pathInArchive, err)
 	}
+	defer fileInArchive.Close()
 
-	out, err := os.OpenFile(ops.LocalPath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0744)
+	out, err := os.CreateTemp(filepath.Dir(localPath), filepath.Base(localPath)+".tmp")
 	if err != nil {
-		return fmt.Errorf("failed to create file %s: %w", ops.LocalPath, err)
+		return fmt.Errorf("failed to create file %s: %w", localPath, err)
 	}
-	_, err = io.Copy(out, fileInArchive)
-	if err != nil {
+	tmpPath := out.Name()
+	defer os.Remove(tmpPath)
+
+	if _, err := io.Copy(out, fileInArchive); err != nil {
+		out.Close()
 		return fmt.Errorf("failed to copy body to file: %w", err)
 	}
-	err = out.Close()
-	if err != nil {
+	if err := out.Close(); err != nil {
 		return fmt.Errorf("failed to close file: %w", err)
 	}
+	if err := os.Chmod(tmpPath, 0744); err != nil {
+		return fmt.Errorf("failed to set permissions on %s: %w", localPath, err)
+	}
+	if err := os.Rename(tmpPath, localPath); err != nil {
+		return fmt.Errorf("failed to move plugin binary to %s: %w", localPath, err)
+	}
 	return nil
 }
 
@@ -284,16 +303,12 @@ func doDownloadPluginFromGithub(ctx context.Context, logger zerolog.Logger, loca
 		return fmt.Errorf("failed to get plugin url: %w", err)
 	}
 	logger.Debug().Msg(fmt.Sprintf("Downloading %s", downloadURL))
+	defer os.Remove(pluginZipPath)
+
 	if _, err := downloadFile(ctx, pluginZipPath, downloadURL, dops); err != nil {
 		return fmt.Errorf("failed to download plugin: %w", err)
 	}
 
-	archive, err := zip.OpenReader(pluginZipPath)
-	if err != nil {
-		return fmt.Errorf("failed to open plugin archive: %w", err)
-	}
-	defer archive.Close()
-
 	var pathInArchive string
 	switch {
 	case strings.HasPrefix(downloadURL, "https://github.com/cloudquery/cloudquery/releases/download/plugins-plugin"):
@@ -312,24 +327,7 @@ func doDownloadPluginFromGithub(ctx context.Context, logger zerolog.Logger, loca
 		return fmt.Errorf("unknown GitHub %s", downloadURL)
 	}
 
-	pathInArchive = WithBinarySuffix(pathInArchive)
-	fileInArchive, err := archive.Open(pathInArchive)
-	if err != nil {
-		return fmt.Errorf("failed to open plugin archive plugins/source/%s: %w", name, err)
-	}
-	out, err := os.OpenFile(localPath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0744)
-	if err != nil {
-		return fmt.Errorf("failed to create file %s: %w", localPath, err)
-	}
-	_, err = io.Copy(out, fileInArchive)
-	if err != nil {
-		return fmt.Errorf("failed to copy body to file: %w", err)
-	}
-	err = out.Close()
-	if err != nil {
-		return fmt.Errorf("failed to close file: %w", err)
-	}
-	return nil
+	return extractPluginBinary(pluginZipPath, WithBinarySuffix(pathInArchive), localPath)
 }
 
 func downloadFile(ctx context.Context, localPath string, downloadURL string, dops DownloaderOptions) (string, error) {

managedplugin/extract_test.go

@@ -0,0 +1,66 @@
+package managedplugin
+
+import (
+	"archive/zip"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func writeTestArchive(t *testing.T, dir, entry string, contents []byte) string {
+	t.Helper()
+
+	archivePath := filepath.Join(dir, "plugin.zip")
+	f, err := os.Create(archivePath)
+	require.NoError(t, err)
+
+	w := zip.NewWriter(f)
+	entryWriter, err := w.Create(entry)
+	require.NoError(t, err)
+	_, err = entryWriter.Write(contents)
+	require.NoError(t, err)
+	require.NoError(t, w.Close())
+	require.NoError(t, f.Close())
+
+	return archivePath
+}
+
+func TestExtractPluginBinary(t *testing.T) {
+	dir := t.TempDir()
+	binary := []byte("plugin-binary")
+	archivePath := writeTestArchive(t, dir, "plugin-aws-v1.0.0-linux-amd64", binary)
+	localPath := filepath.Join(dir, "aws")
+
+	require.NoError(t, extractPluginBinary(archivePath, "plugin-aws-v1.0.0-linux-amd64", localPath))
+
+	got, err := os.ReadFile(localPath)
+	require.NoError(t, err)
+	require.Equal(t, binary, got)
+
+	info, err := os.Stat(localPath)
+	require.NoError(t, err)
+	require.Equal(t, os.FileMode(0744), info.Mode().Perm())
+}
+
+// TestExtractPluginBinaryLeavesNoPartialFile guards the caching path: a failed
+// extraction that left bytes at localPath would make the next run skip the
+// download and exec a truncated binary.
+func TestExtractPluginBinaryLeavesNoPartialFile(t *testing.T) {
+	dir := t.TempDir()
+	archivePath := writeTestArchive(t, dir, "plugin-aws-v1.0.0-linux-amd64", []byte("plugin-binary"))
+	localPath := filepath.Join(dir, "aws")
+
+	err := extractPluginBinary(archivePath, "plugin-aws-v9.9.9-linux-amd64", localPath)
+	require.Error(t, err)
+
+	_, statErr := os.Stat(localPath)
+	require.ErrorIs(t, statErr, os.ErrNotExist)
+
+	entries, err := os.ReadDir(dir)
+	require.NoError(t, err)
+	for _, entry := range entries {
+		require.NotContains(t, entry.Name(), ".tmp", "the temporary file must not survive a failed extraction")
+	}
+}