From 43e03be7fd27faa2faf80b20a2d73acaadc18c9a Mon Sep 17 00:00:00 2001
From: dbw7 <danial.bekhit@suse.com>
Date: Thu, 1 Aug 2024 09:23:06 -0400
Subject: [PATCH] Combustion Cleanup (#514)

* cleanup combustion
---
 RELEASE_NOTES.md                              |  1 +
 pkg/combustion/cleanup.go                     | 37 +++++++++++++
 pkg/combustion/cleanup_test.go                | 52 +++++++++++++++++++
 pkg/combustion/combustion.go                  |  7 +++
 .../templates/cleanup-combustion.sh           |  8 +++
 5 files changed, 105 insertions(+)
 create mode 100644 pkg/combustion/cleanup.go
 create mode 100644 pkg/combustion/cleanup_test.go
 create mode 100755 pkg/combustion/templates/cleanup-combustion.sh

diff --git a/RELEASE_NOTES.md b/RELEASE_NOTES.md
index 86837ce2..24d4c8c0 100644
--- a/RELEASE_NOTES.md
+++ b/RELEASE_NOTES.md
@@ -11,6 +11,7 @@
 * Added ability to build aarch64 images on an aarch64 host machine
 * Added caching for container images
 * Added built image name output to build command
+* Leftover combustion artifacts are now removed on first boot
 
 ## API
 
diff --git a/pkg/combustion/cleanup.go b/pkg/combustion/cleanup.go
new file mode 100644
index 00000000..1d72a931
--- /dev/null
+++ b/pkg/combustion/cleanup.go
@@ -0,0 +1,37 @@
+package combustion
+
+import (
+	_ "embed"
+	"fmt"
+	"os"
+	"path/filepath"
+
+	"github.com/suse-edge/edge-image-builder/pkg/fileio"
+	"github.com/suse-edge/edge-image-builder/pkg/image"
+	"github.com/suse-edge/edge-image-builder/pkg/log"
+	"go.uber.org/zap"
+)
+
+const (
+	cleanupScriptName    = "cleanup-combustion.sh"
+	cleanupComponentName = "cleanup"
+)
+
+//go:embed templates/cleanup-combustion.sh
+var cleanupScript string
+
+func configureCleanup(ctx *image.Context) ([]string, error) {
+	if ctx.ImageDefinition.Image.ImageType != image.TypeRAW {
+		log.AuditComponentSkipped(cleanupComponentName)
+		zap.S().Info("skipping cleanup component, image type is not raw")
+		return nil, nil
+	}
+
+	cleanupScriptFilename := filepath.Join(ctx.CombustionDir, cleanupScriptName)
+	if err := os.WriteFile(cleanupScriptFilename, []byte(cleanupScript), fileio.ExecutablePerms); err != nil {
+		return nil, fmt.Errorf("writing cleanup files script %s: %w", cleanupScriptName, err)
+	}
+
+	log.AuditComponentSuccessful(cleanupComponentName)
+	return []string{cleanupScriptName}, nil
+}
diff --git a/pkg/combustion/cleanup_test.go b/pkg/combustion/cleanup_test.go
new file mode 100644
index 00000000..80d43e92
--- /dev/null
+++ b/pkg/combustion/cleanup_test.go
@@ -0,0 +1,52 @@
+package combustion
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/suse-edge/edge-image-builder/pkg/image"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestConfigureCleanupRaw(t *testing.T) {
+	// Setup
+	ctx, teardown := setupContext(t)
+	defer teardown()
+	ctx.ImageDefinition.Image.ImageType = image.TypeRAW
+
+	// Test
+	scriptNames, err := configureCleanup(ctx)
+
+	// Verify
+	require.NoError(t, err)
+
+	assert.Equal(t, []string{cleanupScriptName}, scriptNames)
+
+	// -- Combustion Script
+	expectedCombustionScript := filepath.Join(ctx.CombustionDir, cleanupScriptName)
+	contents, err := os.ReadFile(expectedCombustionScript)
+	require.NoError(t, err)
+	assert.Contains(t, string(contents), "rm -r /artefacts")
+}
+
+func TestConfigureCleanupISO(t *testing.T) {
+	// Setup
+	ctx, teardown := setupContext(t)
+	defer teardown()
+	ctx.ImageDefinition.Image.ImageType = image.TypeISO
+
+	// Test
+	scriptNames, err := configureCleanup(ctx)
+
+	// Verify
+	require.NoError(t, err)
+
+	assert.NotEqual(t, []string{cleanupScriptName}, scriptNames)
+
+	// -- Combustion Script
+	expectedCombustionScript := filepath.Join(ctx.CombustionDir, cleanupScriptName)
+	assert.NoFileExists(t, expectedCombustionScript)
+}
diff --git a/pkg/combustion/combustion.go b/pkg/combustion/combustion.go
index 035d17fc..0a139cfc 100644
--- a/pkg/combustion/combustion.go
+++ b/pkg/combustion/combustion.go
@@ -163,6 +163,13 @@ func (c *Combustion) Configure(ctx *image.Context) error {
 		combustionScripts = append(combustionScripts, scripts...)
 	}
 
+	// We manually add the cleanup component as to always make sure it is last
+	s, err := configureCleanup(ctx)
+	if err != nil {
+		return fmt.Errorf("configuring cleanup component %q: %w", cleanupComponentName, err)
+	}
+	combustionScripts = append(combustionScripts, s...)
+
 	var networkScript string
 	if isComponentConfigured(ctx, networkConfigDir) {
 		networkScript = networkConfigScriptName
diff --git a/pkg/combustion/templates/cleanup-combustion.sh b/pkg/combustion/templates/cleanup-combustion.sh
new file mode 100755
index 00000000..49cc3444
--- /dev/null
+++ b/pkg/combustion/templates/cleanup-combustion.sh
@@ -0,0 +1,8 @@
+#!/bin/bash
+set -euo pipefail
+
+rm -r /combustion
+
+if test -d /artefacts; then
+  rm -r /artefacts
+fi
\ No newline at end of file