Merge pull request #69 from OpenCHAMI/synackd/yaml-marshallers

feat(cistore): add `MarshalYAML()` and `UnmarshalYAML()` to `CloudConfigFile`

Author: synackd

go.mod

@@ -12,6 +12,7 @@ require (
 	github.com/rs/zerolog v1.33.0
 	github.com/stretchr/testify v1.10.0
 	github.com/swaggo/swag v1.16.4
+	gopkg.in/yaml.v3 v3.0.1
 )
 
 require (
@@ -40,7 +41,6 @@ require (
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
 	github.com/segmentio/asm v1.2.0 // indirect
 	golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d // indirect
-	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
 
 require (

pkg/cistore/models.go

@@ -3,6 +3,9 @@ package cistore
 import (
 	"encoding/json"
 	"errors"
+	"fmt"
+
+	"gopkg.in/yaml.v3"
 
 	base "github.com/Cray-HPE/hms-base"
 )
@@ -65,12 +68,12 @@ type ClusterDefaults struct {
 }
 
 type CloudConfigFile struct {
-	Content  []byte `json:"content" swaggertype:"string" example:"IyMgdGVtcGxhdGU6IGppbmphCiNjbG91ZC1jb25maWcKbWVyZ2VfaG93OgotIG5hbWU6IGxpc3QKICBzZXR0aW5nczogW2FwcGVuZF0KLSBuYW1lOiBkaWN0CiAgc2V0dGluZ3M6IFtub19yZXBsYWNlLCByZWN1cnNlX2xpc3RdCnVzZXJzOgogIC0gbmFtZTogcm9vdAogICAgc3NoX2F1dGhvcml6ZWRfa2V5czoge3sgZHMubWV0YV9kYXRhLmluc3RhbmNlX2RhdGEudjEucHVibGljX2tleXMgfX0KZGlzYWJsZV9yb290OiBmYWxzZQo=" description:"Cloud-Init configuration content whose encoding depends on the value of 'encoding'"`
-	Name     string `json:"filename"`
-	Encoding string `json:"encoding,omitempty" enums:"base64,plain"`
+	Content  []byte `json:"content" yaml:"content" swaggertype:"string" example:"IyMgdGVtcGxhdGU6IGppbmphCiNjbG91ZC1jb25maWcKbWVyZ2VfaG93OgotIG5hbWU6IGxpc3QKICBzZXR0aW5nczogW2FwcGVuZF0KLSBuYW1lOiBkaWN0CiAgc2V0dGluZ3M6IFtub19yZXBsYWNlLCByZWN1cnNlX2xpc3RdCnVzZXJzOgogIC0gbmFtZTogcm9vdAogICAgc3NoX2F1dGhvcml6ZWRfa2V5czoge3sgZHMubWV0YV9kYXRhLmluc3RhbmNlX2RhdGEudjEucHVibGljX2tleXMgfX0KZGlzYWJsZV9yb290OiBmYWxzZQo=" description:"Cloud-Init configuration content whose encoding depends on the value of 'encoding'"`
+	Name     string `json:"filename" yaml:"filename"`
+	Encoding string `json:"encoding,omitempty" yaml:"encoding,omitempty" enums:"base64,plain"`
 }
 
-// Custom unmarshaler for CloudConfigFile
+// Custom JSON unmarshaler for CloudConfigFile
 func (f *CloudConfigFile) UnmarshalJSON(data []byte) error {
 	// Use an auxiliary struct so that:
 	//
@@ -96,11 +99,57 @@ func (f *CloudConfigFile) UnmarshalJSON(data []byte) error {
 	return nil
 }
 
-// Custom marshaler for CloudConfigFile
+// Custom YAML unmarshaller for CloudConfigFile. This is needed because the yaml
+// library cannot unmarshal string to []byte like json can, so it needs to be
+// told how to do so.
+func (f *CloudConfigFile) UnmarshalYAML(n *yaml.Node) error {
+	// Use an auxiliary struct so that:
+	//
+	// 1. json.Unmarshal doesn't recurse forever and overflow the stack.
+	// 2. json.Unmarshal doesn't try to base64-decode "content" in the data
+	//    before assigning the bytes to f.Content. Content is unmarshalled
+	//    as a string instead of bytes in order to prevent this. After
+	//    unmarshalling, the string is converted back to bytes and assigned
+	//    to f.Content.
+	//
+	// Interestingly, yaml.Unmarshal will not unmarshal into aux's pointer
+	// to f, so we have to use json.Unmarshal.
+	type Alias CloudConfigFile
+	aux := &struct {
+		Content string `json:"content"`
+		*Alias
+	}{
+		Alias: (*Alias)(f),
+	}
+
+	// Decode YAML document (n) into aux struct, using a map as an
+	// intermediary (since n.Decode() cannot decode into a []byte). We have
+	// to use JSON for the unmarshalling (and, by consequence, the
+	// marshalling) so that f will get written via aux's pointer to it. For
+	// some reason, the yaml unmarshaller will not do that.
+	//
+	// 1. Decode YAML document (n) into map.
+	// 2. JSON marshal map into bytes.
+	// 3. JSON unmarshal bytes into aux struct.
+	// 4. Set f.Content to byte-ified aux.Content.
+	var mAux map[string]interface{}
+	if err := n.Decode(&mAux); err != nil {
+		return err
+	}
+	t, err := json.Marshal(mAux)
+	if err != nil {
+		return err
+	}
+	if err := json.Unmarshal(t, &aux); err != nil {
+		return err
+	}
+	f.Content = []byte(aux.Content)
+
+	return nil
+}
+
+// Custom JSON marshaler for CloudConfigFile
 func (f CloudConfigFile) MarshalJSON() ([]byte, error) {
-	// Use temporary struct to marshal so json.Marshal doesn't recurse
-	// indefinitely. Also to convert Content from bytes to string so
-	// json.Marshal doesn't try to base64 encode the bytes.
 	// Use an auxiliary struct so that:
 	//
 	// 1. json.Marshal doesn't recurse forever and overflow the stack.
@@ -118,3 +167,51 @@ func (f CloudConfigFile) MarshalJSON() ([]byte, error) {
 	aux.Content = string(f.Content)
 	return json.Marshal(aux)
 }
+
+// Custom YAML marshaler for CloudConfigFile. This is needed because the yaml
+// library will not marshal []byte into string like json will, so it needs to be
+// told how to do so.
+func (f CloudConfigFile) MarshalYAML() (interface{}, error) {
+	// Use an auxiliary struct so that:
+	//
+	// 1. yaml.Marshal doesn't recurse forever and overflow the stack.
+	// 2. yaml.Marshal doesn't try to base64-encode f.Content. f.Content is
+	//    converted from bytes to a string and then assigned to aux.Content
+	//    to prevent this. Then, aux gets marshalled instead of f.
+	//
+	// aux is set to the values of f, but has its own string Content, set to
+	// the stringified f.Content.
+	type Alias CloudConfigFile
+	aux := &struct {
+		Content string `yaml:"content"`
+		Alias
+	}{
+		Content: string(f.Content),
+		Alias:   (Alias)(f),
+	}
+
+	// Convert aux into map, which has an "alias" key containing the
+	// "content" that will be set to the actual string value. The map that
+	// is mapped to "alias" is what is returned.
+	//
+	// 1. YAML marshal aux to bytes.
+	// 2. YAML unmarshal bytes into map.
+	// 3. Set content in "alias" to aux.Content.
+	// 4. Return "alias" map.
+	t, err := yaml.Marshal(aux)
+	if err != nil {
+		return nil, err
+	}
+	var nMap map[string]interface{}
+	if err := yaml.Unmarshal(t, &nMap); err != nil {
+		return nil, err
+	}
+	switch c := (nMap["alias"]).(type) {
+	case map[string]interface{}:
+		c["content"] = aux.Content
+	default:
+		return nil, fmt.Errorf("cloud config file in map is unknown type: wanted=(map[string]interface{}) actual=(%v)", c)
+	}
+
+	return nMap["alias"], nil
+}

pkg/cistore/models_test.go

@@ -3,6 +3,7 @@ package cistore
 import (
 	"encoding/base64"
 	"encoding/json"
+	"gopkg.in/yaml.v3"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -23,6 +24,19 @@ func TestCloudConfigFile_UnmarshalJSON_Plain(t *testing.T) {
 	assert.Equal(t, []byte("#cloud-config\nusers:\n  - name: test"), f.Content)
 }
 
+func TestCloudConfigFile_UnmarshalYAML_Plain(t *testing.T) {
+	yamlData := []byte(`filename: myconfig.yaml
+encoding: plain
+content: "#cloud-config\nusers:\n  - name: test"`)
+
+	var f CloudConfigFile
+	err := yaml.Unmarshal(yamlData, &f)
+	assert.NoError(t, err)
+	assert.Equal(t, "myconfig.yaml", f.Name)
+	assert.Equal(t, "plain", f.Encoding)
+	assert.Equal(t, []byte("#cloud-config\nusers:\n  - name: test"), f.Content)
+}
+
 func TestCloudConfigFile_UnmarshalJSON_Base64(t *testing.T) {
 	encodedContent := base64.StdEncoding.EncodeToString([]byte("#cloud-config\nusers:\n  - name: test"))
 	jsonData := []byte(`{
@@ -40,6 +54,20 @@ func TestCloudConfigFile_UnmarshalJSON_Base64(t *testing.T) {
 	assert.Equal(t, []byte(encodedContent), f.Content)
 }
 
+func TestCloudConfigFile_UnmarshalYAML_Base64(t *testing.T) {
+	encodedContent := base64.StdEncoding.EncodeToString([]byte("#cloud-config\nusers:\n  - name: test"))
+	yamlData := []byte(`filename: myconfig.yaml
+encoding: base64
+content: ` + encodedContent)
+
+	var f CloudConfigFile
+	err := yaml.Unmarshal(yamlData, &f)
+	assert.NoError(t, err)
+	assert.Equal(t, "myconfig.yaml", f.Name)
+	assert.Equal(t, "base64", f.Encoding)
+	assert.Equal(t, []byte(encodedContent), f.Content)
+}
+
 func TestCloudConfigFile_MarshalJSON_Plain(t *testing.T) {
 	f := CloudConfigFile{
 		Name:     "plainconfig.yaml",
@@ -58,6 +86,24 @@ func TestCloudConfigFile_MarshalJSON_Plain(t *testing.T) {
 	assert.Equal(t, "#cloud-config\nusers:\n  - name: test", out["content"])
 }
 
+func TestCloudConfigFile_MarshalYAML_Plain(t *testing.T) {
+	f := CloudConfigFile{
+		Name:     "plainconfig.yaml",
+		Encoding: "plain",
+		Content:  []byte("#cloud-config\nusers:\n  - name: test"),
+	}
+
+	data, err := yaml.Marshal(f)
+	assert.NoError(t, err)
+
+	var out map[string]interface{}
+	err = yaml.Unmarshal(data, &out)
+	assert.NoError(t, err)
+	assert.Equal(t, "plainconfig.yaml", out["filename"])
+	assert.Equal(t, "plain", out["encoding"])
+	assert.Equal(t, "#cloud-config\nusers:\n  - name: test", out["content"])
+}
+
 func TestCloudConfigFile_MarshalJSON_Base64(t *testing.T) {
 	originalConfig := "#cloud-config\nusers:\n  - name: test"
 	b64Config := base64.StdEncoding.EncodeToString([]byte(originalConfig))
@@ -77,3 +123,23 @@ func TestCloudConfigFile_MarshalJSON_Base64(t *testing.T) {
 	assert.Equal(t, "base64", out["encoding"])
 	assert.Equal(t, b64Config, out["content"])
 }
+
+func TestCloudConfigFile_MarshalYAML_Base64(t *testing.T) {
+	originalConfig := "#cloud-config\nusers:\n  - name: test"
+	b64Config := base64.StdEncoding.EncodeToString([]byte(originalConfig))
+	f := CloudConfigFile{
+		Name:     "encodedconfig.yaml",
+		Encoding: "base64",
+		Content:  []byte(b64Config),
+	}
+
+	data, err := yaml.Marshal(f)
+	assert.NoError(t, err)
+
+	var out map[string]interface{}
+	err = yaml.Unmarshal(data, &out)
+	assert.NoError(t, err)
+	assert.Equal(t, "encodedconfig.yaml", out["filename"])
+	assert.Equal(t, "base64", out["encoding"])
+	assert.Equal(t, b64Config, out["content"])
+}