From 3a502e32f399591614a041f1c96b9c5f7138707f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Edu=20G=C3=B3mez=20Escandell?=
 <edu.gomez.escandell@canonical.com>
Date: Tue, 23 Jan 2024 14:23:02 +0100
Subject: [PATCH] Have Landscape write and remove cloud-init user data files

---
 .../proservices/landscape/executor.go         | 43 ++++++++++++++-----
 1 file changed, 32 insertions(+), 11 deletions(-)

diff --git a/windows-agent/internal/proservices/landscape/executor.go b/windows-agent/internal/proservices/landscape/executor.go
index f68a0b7e1..ff9d34fc9 100644
--- a/windows-agent/internal/proservices/landscape/executor.go
+++ b/windows-agent/internal/proservices/landscape/executor.go
@@ -7,6 +7,7 @@ import (
 	"os/user"
 
 	landscapeapi "github.com/canonical/landscape-hostagent-api"
+	"github.com/canonical/ubuntu-pro-for-windows/windows-agent/internal/cloudinit"
 	log "github.com/canonical/ubuntu-pro-for-windows/windows-agent/internal/grpc/logstreamer"
 	"github.com/canonical/ubuntu-pro-for-windows/windows-agent/internal/proservices/landscape/distroinstall"
 	"github.com/ubuntu/decorate"
@@ -100,9 +101,6 @@ func (e executor) stop(ctx context.Context, cmd *landscapeapi.Command_Stop) (err
 
 func (e executor) install(ctx context.Context, cmd *landscapeapi.Command_Install) (err error) {
 	log.Debugf(ctx, "Landscape: received command Install. Target: %s", cmd.GetId())
-	if cmd.GetCloudinit() != "" {
-		return fmt.Errorf("Cloud Init support is not yet available")
-	}
 
 	distro := gowsl.NewDistro(ctx, cmd.GetId())
 	if registered, err := distro.IsRegistered(); err != nil {
@@ -111,20 +109,28 @@ func (e executor) install(ctx context.Context, cmd *landscapeapi.Command_Install
 		return errors.New("already installed")
 	}
 
+	if err := cloudinit.WriteDistroData(distro.Name(), "landscape", cmd.GetCloudinit()); err != nil {
+		return fmt.Errorf("skipped install: %v", err)
+	}
+
+	// Avoid dangling cloud-init file
+	defer doOnError(&err, func() {
+		if err := cloudinit.RemoveDistroData(distro.Name(), "landscape"); err != nil {
+			log.Warningf(ctx, "Landscape Install: distro %q: could not remove cloud-init file after failing to install: %v", distro.Name(), err)
+		}
+	})
+
 	if err := gowsl.Install(ctx, distro.Name()); err != nil {
 		return err
 	}
 
-	defer func() {
-		if err == nil {
-			return
-		}
-		// Avoid error states by cleaning up on error
+	// Avoid dangling distro
+	defer doOnError(&err, func() {
 		err := distro.Uninstall(ctx)
 		if err != nil {
-			log.Infof(ctx, "Landscape Install: failed to clean up %q after failed Install: %v", distro.Name(), err)
+			log.Infof(ctx, "Landscape Install: distro %q: failed to uninstall after failed Install: %v", distro.Name(), err)
 		}
-	}()
+	})
 
 	if err := distroinstall.InstallFromExecutable(ctx, distro); err != nil {
 		return err
@@ -153,6 +159,13 @@ func (e executor) install(ctx context.Context, cmd *landscapeapi.Command_Install
 	return nil
 }
 
+// doOnError to avod verbose deferred cleanups.
+func doOnError(err *error, f func()) {
+	if *err != nil {
+		f()
+	}
+}
+
 func (e executor) uninstall(ctx context.Context, cmd *landscapeapi.Command_Uninstall) (err error) {
 	log.Debugf(ctx, "Landscape: received command Uninstall. Target: %s", cmd.GetId())
 	d, ok := e.database().Get(cmd.GetId())
@@ -160,7 +173,15 @@ func (e executor) uninstall(ctx context.Context, cmd *landscapeapi.Command_Unins
 		return fmt.Errorf("distro %q not in database", cmd.GetId())
 	}
 
-	return d.Uninstall(ctx)
+	if err := d.Uninstall(ctx); err != nil {
+		return err
+	}
+
+	if err := cloudinit.RemoveDistroData(d.Name(), "landscape"); err != nil {
+		log.Warningf(ctx, "Landscape uninstall: distro %q: %v", d.Name(), err)
+	}
+
+	return nil
 }
 
 func (e executor) setDefault(ctx context.Context, cmd *landscapeapi.Command_SetDefault) error {