From 00784a072a3c843c73b1ccbf45281fd77af5112b Mon Sep 17 00:00:00 2001
From: Roman Ryzhyi <ryzhyi.r@ajax.systems>
Date: Sun, 12 Nov 2023 14:00:09 +0200
Subject: [PATCH] Add source image ID and region attributes

---
 builder/hcloud/artifact.go      | 38 +++++++++++++++++++++++-----
 builder/hcloud/artifact_test.go | 44 +++++++++++++++++++++++++++++++++
 builder/hcloud/builder.go       |  7 +++++-
 3 files changed, 82 insertions(+), 7 deletions(-)

diff --git a/builder/hcloud/artifact.go b/builder/hcloud/artifact.go
index f85c7fed..d7f54d8c 100644
--- a/builder/hcloud/artifact.go
+++ b/builder/hcloud/artifact.go
@@ -47,16 +47,42 @@ func (a *Artifact) String() string {
 
 func (a *Artifact) State(name string) interface{} {
 	if name == registryimage.ArtifactStateURI {
-		img, err := registryimage.FromArtifact(a, registryimage.WithProvider("hetznercloud"))
-		if err != nil {
-			log.Printf("[DEBUG] error encountered when creating a registry image %v", err)
-			return nil
-		}
-		return img
+		return a.stateHCPPackerRegistryMetadata()
 	}
 	return a.StateData[name]
 }
 
+func (a *Artifact) stateHCPPackerRegistryMetadata() interface{} {
+	// create labels map
+	labels := make(map[string]string)
+	sourceImage, ok := a.StateData["source_image"].(string)
+	if ok {
+		labels["source_image"] = sourceImage
+	}
+	// get and set region from stateData into labels
+	region, ok := a.StateData["region"].(string)
+	if ok {
+		labels["region"] = region
+	}
+	// get and set server_type from stateData into labels
+	serverType, ok := a.StateData["server_type"].(string)
+	if ok {
+		labels["server_type"] = serverType
+	}
+
+	image, err := registryimage.FromArtifact(a,
+		registryimage.WithProvider("hetznercloud"),
+		registryimage.WithID(strconv.FormatInt(a.snapshotId, 10)),
+		registryimage.WithSourceID(sourceImage),
+		registryimage.WithRegion(region))
+	if err != nil {
+		log.Printf("[DEBUG] error encountered when creating registry image %s", err)
+		return nil
+	}
+	image.Labels = labels
+	return image
+}
+
 func (a *Artifact) Destroy() error {
 	log.Printf("Destroying image: %d (%s)", a.snapshotId, a.snapshotName)
 	_, err := a.hcloudClient.Image.Delete(context.TODO(), &hcloud.Image{ID: a.snapshotId})
diff --git a/builder/hcloud/artifact_test.go b/builder/hcloud/artifact_test.go
index 25870859..d9b3ab63 100644
--- a/builder/hcloud/artifact_test.go
+++ b/builder/hcloud/artifact_test.go
@@ -4,9 +4,12 @@
 package hcloud
 
 import (
+	"reflect"
 	"testing"
 
 	packersdk "github.com/hashicorp/packer-plugin-sdk/packer"
+	registryimage "github.com/hashicorp/packer-plugin-sdk/packer/registry/image"
+	"github.com/mitchellh/mapstructure"
 )
 
 func TestArtifact_Impl(t *testing.T) {
@@ -58,3 +61,44 @@ func TestArtifactState_StateData(t *testing.T) {
 		t.Fatalf("Bad: State should be nil for nil StateData")
 	}
 }
+
+func TestArtifactState_hcpPackerRegistryMetadata(t *testing.T) {
+	region := "nbg1"
+	artifact := &Artifact{
+		snapshotId:   1111111111,
+		snapshotName: "test-image",
+		StateData: map[string]interface{}{
+			"source_image": "ubuntu-22.04",
+			"region":       region,
+			"server_type":  "cx11",
+		},
+	}
+	// result should contain "something"
+	result := artifact.State(registryimage.ArtifactStateURI)
+	if result == nil {
+		t.Fatalf("Bad: HCP Packer registry image data was nil")
+	}
+
+	// check for proper decoding of result into slice of registryimage.Image
+	var image registryimage.Image
+	err := mapstructure.Decode(result, &image)
+	if err != nil {
+		t.Errorf("Bad: unexpected error when trying to decode state into registryimage.Image %v", err)
+	}
+
+	// check that all properties of the images were set correctly
+	expected := registryimage.Image{
+		ImageID:        "1111111111",
+		ProviderName:   "hetznercloud",
+		ProviderRegion: region,
+		SourceImageID:  "ubuntu-22.04",
+		Labels: map[string]string{
+			"source_image": "ubuntu-22.04",
+			"region":       region,
+			"server_type":  "cx11",
+		},
+	}
+	if !reflect.DeepEqual(image, expected) {
+		t.Fatalf("Bad: expected %#v got %#v", expected, image)
+	}
+}
diff --git a/builder/hcloud/builder.go b/builder/hcloud/builder.go
index ff37bf03..5667cd07 100644
--- a/builder/hcloud/builder.go
+++ b/builder/hcloud/builder.go
@@ -98,7 +98,12 @@ func (b *Builder) Run(ctx context.Context, ui packersdk.Ui, hook packersdk.Hook)
 		snapshotName: state.Get("snapshot_name").(string),
 		snapshotId:   state.Get("snapshot_id").(int64),
 		hcloudClient: b.hcloudClient,
-		StateData:    map[string]interface{}{"generated_data": state.Get("generated_data")},
+		StateData: map[string]interface{}{
+			"generated_data": state.Get("generated_data"),
+			"source_image":   b.config.Image,
+			"region":         b.config.Location,
+			"server_type":    b.config.ServerType,
+		},
 	}
 
 	return artifact, nil