From 8911486fcad50b151399e4652c7c62f2e037e7dd Mon Sep 17 00:00:00 2001
From: Utku Ozdemir <utku.ozdemir@siderolabs.com>
Date: Tue, 10 Sep 2024 14:25:03 +0200
Subject: [PATCH] feat: respect custom JSON unmarshaler logic in protobuf
 resource specs

If a protobuf spec implements `json.Unmarshaler`, use that function to unmarshal it from JSON, as the default `protojson.UnmarshalOptions` ignores it.

Add a test for both JSON and YAML custom unmarshaling logic.

Co-authored-by: Artem Chernyshev <artem.chernyshev@talos-systems.com>
Signed-off-by: Utku Ozdemir <utku.ozdemir@siderolabs.com>
---
 pkg/resource/protobuf/spec.go                |  4 +
 pkg/resource/protobuf/spec_unmarshal_test.go | 81 ++++++++++++++++++++
 2 files changed, 85 insertions(+)
 create mode 100644 pkg/resource/protobuf/spec_unmarshal_test.go

diff --git a/pkg/resource/protobuf/spec.go b/pkg/resource/protobuf/spec.go
index dc2a4962..a1d43bf5 100644
--- a/pkg/resource/protobuf/spec.go
+++ b/pkg/resource/protobuf/spec.go
@@ -70,6 +70,10 @@ func (spec *ResourceSpec[T, S]) UnmarshalYAML(node *yaml.Node) error {
 func (spec *ResourceSpec[T, S]) UnmarshalJSON(bytes []byte) error {
 	spec.Value = new(T)
 
+	if unmarshaler, ok := any(spec.Value).(json.Unmarshaler); ok {
+		return unmarshaler.UnmarshalJSON(bytes)
+	}
+
 	opts := protojson.UnmarshalOptions{}
 
 	return opts.Unmarshal(bytes, spec.Value)
diff --git a/pkg/resource/protobuf/spec_unmarshal_test.go b/pkg/resource/protobuf/spec_unmarshal_test.go
new file mode 100644
index 00000000..de6a1116
--- /dev/null
+++ b/pkg/resource/protobuf/spec_unmarshal_test.go
@@ -0,0 +1,81 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+package protobuf_test
+
+import (
+	"encoding/json"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/reflect/protoreflect"
+	"gopkg.in/yaml.v3"
+
+	"github.com/cosi-project/runtime/pkg/resource/protobuf"
+)
+
+type rawSpec struct {
+	Str string
+	Num int
+}
+
+type customUnmarshalerSpec struct {
+	Str string
+	Num int
+}
+
+func (spec *customUnmarshalerSpec) ProtoReflect() protoreflect.Message { return nil }
+
+// UnmarshalJSON uppercases the string and doubles the number.
+func (spec *customUnmarshalerSpec) UnmarshalJSON(data []byte) error {
+	var raw rawSpec
+
+	err := json.Unmarshal(data, &raw)
+	if err != nil {
+		return err
+	}
+
+	spec.Str = strings.ToUpper(raw.Str)
+	spec.Num = raw.Num * 2
+
+	return nil
+}
+
+// UnmarshalYAML lowercases the string and halves the number.
+func (spec *customUnmarshalerSpec) UnmarshalYAML(node *yaml.Node) error {
+	var raw rawSpec
+
+	err := node.Decode(&raw)
+	if err != nil {
+		return err
+	}
+
+	spec.Str = strings.ToLower(raw.Str)
+	spec.Num = raw.Num / 2
+
+	return nil
+}
+
+func TestCustomJSONUnmarshal(t *testing.T) {
+	spec := protobuf.NewResourceSpec(&customUnmarshalerSpec{})
+
+	err := json.Unmarshal([]byte(`{"str":"aaaa","num":2222}`), &spec)
+	require.NoError(t, err)
+
+	assert.Equal(t, "AAAA", spec.Value.Str)
+	assert.Equal(t, 4444, spec.Value.Num)
+}
+
+func TestCustomYAMLUnmarshal(t *testing.T) {
+	spec := protobuf.NewResourceSpec(&customUnmarshalerSpec{})
+
+	err := yaml.Unmarshal([]byte(`str: AAAA
+num: 2222`), &spec)
+	require.NoError(t, err)
+
+	assert.Equal(t, "aaaa", spec.Value.Str)
+	assert.Equal(t, 1111, spec.Value.Num)
+}