Preserve extension fields when marshalling/unmarshalling

This makes it so external callers don't need to handle this themselves.

Signed-off-by: Brian Goff <[email protected]>

load.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 
 	"github.com/goccy/go-yaml"
+	"github.com/goccy/go-yaml/parser"
 	"github.com/moby/buildkit/frontend/dockerfile/shell"
 	"github.com/pkg/errors"
 	"golang.org/x/exp/maps"
@@ -210,36 +211,93 @@ func (s *Spec) SubstituteArgs(env map[string]string, opts ...SubstituteOpt) erro
 func LoadSpec(dt []byte) (*Spec, error) {
 	var spec Spec
 
-	dt, err := stripXFields(dt)
-	if err != nil {
-		return nil, fmt.Errorf("error stripping x-fields: %w", err)
-	}
-
 	if err := yaml.UnmarshalWithOptions(dt, &spec, yaml.Strict()); err != nil {
-		return nil, fmt.Errorf("error unmarshalling spec: %w", err)
+		return nil, errors.Wrap(err, "error unmarshalling spec")
 	}
 
 	if err := spec.Validate(); err != nil {
 		return nil, err
 	}
-	spec.FillDefaults()
 
+	spec.FillDefaults()
 	return &spec, nil
 }
 
-func stripXFields(dt []byte) ([]byte, error) {
+func (s *Spec) UnmarshalYAML(f func(interface{}) error) error {
 	var obj map[string]interface{}
-	if err := yaml.Unmarshal(dt, &obj); err != nil {
-		return nil, fmt.Errorf("error unmarshalling spec: %w", err)
+
+	if err := f(&obj); err != nil {
+		return err
+	}
+
+	var ext map[string]interface{}
+
+	addExt := func(k string, v interface{}) {
+		if ext == nil {
+			ext = make(map[string]interface{})
+		}
+		ext[k] = v
 	}
 
-	for k := range obj {
-		if strings.HasPrefix(k, "x-") || strings.HasPrefix(k, "X-") {
-			delete(obj, k)
+	for k, v := range obj {
+		if !strings.HasPrefix(k, "x-") && !strings.HasPrefix(k, "X-") {
+			continue
 		}
+		addExt(k, v)
+		delete(obj, k)
+	}
+
+	type internalSpec Spec
+	var s2 internalSpec
+
+	dt, err := yaml.Marshal(obj)
+	if err != nil {
+		return err
+	}
+
+	parsed, err := parser.ParseBytes(dt, 0)
+	if err != nil {
+		return err
+	}
+
+	if err := yaml.UnmarshalWithOptions([]byte(parsed.Docs[0].String()), &s2, yaml.Strict()); err != nil {
+		return err
+	}
+
+	dt, err = yaml.Marshal(ext)
+	if err != nil {
+		return err
+	}
+
+	parsedExt, err := parser.ParseBytes(dt, 0)
+	if err != nil {
+		return errors.Wrap(err, "error parsing extension")
+	}
+
+	*s = Spec(s2)
+	s.ext = ext
+	s.extRaw = parsedExt
+
+	return nil
+}
+
+func (s Spec) MarshalYAML() ([]byte, error) {
+	// We need to define a new type to avoid infinite recursion of MarshalYAML.
+	type internalSpec Spec
+
+	if s.ext == nil {
+		return yaml.Marshal(internalSpec(s))
+	}
+
+	type specExt struct {
+		internalSpec `yaml:",inline"`
+		Ext          map[string]interface{} `yaml:",omitempty,inline"`
 	}
 
-	return yaml.Marshal(obj)
+	return yaml.Marshal(specExt{
+		internalSpec: internalSpec(s),
+		Ext:          s.ext,
+	})
 }
 
 func (s *BuildStep) processBuildArgs(lex *shell.Lex, args map[string]string, allowArg func(string) bool) error {

load_test.go

@@ -10,6 +10,7 @@ import (
 	"reflect"
 	"testing"
 
+	"github.com/goccy/go-yaml"
 	"github.com/moby/buildkit/frontend/dockerui"
 	"gotest.tools/v3/assert"
 	"gotest.tools/v3/assert/cmp"
@@ -1160,3 +1161,52 @@ func TestImage_validate(t *testing.T) {
 		})
 	}
 }
+
+func checkExt[T any](t *testing.T, spec Spec, key string, expect T) {
+	t.Helper()
+
+	var actual T
+	err := spec.Ext(key, &actual)
+	assert.NilError(t, err)
+	assert.Check(t, cmp.DeepEqual(actual, expect))
+}
+
+func TestExtensionFieldMarshalUnmarshal(t *testing.T) {
+	dt := []byte(`
+name: test
+x-hello: world
+x-foo:
+- bar
+- baz
+X-capitalized: world2
+`)
+
+	var spec Spec
+	err := yaml.Unmarshal(dt, &spec)
+	assert.NilError(t, err)
+
+	assert.Check(t, cmp.Equal(spec.Name, "test"), spec)
+	checkExt(t, spec, "hello", "world")
+	checkExt(t, spec, "x-hello", "world")
+	checkExt(t, spec, "foo", []string{"bar", "baz"})
+	checkExt(t, spec, "x-foo", []string{"bar", "baz"})
+	checkExt(t, spec, "capitalized", "world2")
+	checkExt(t, spec, "X-capitalized", "world2")
+
+	// marshal and unmarshal to ensure the extension fields are preserved
+
+	dt, err = yaml.Marshal(spec)
+	assert.NilError(t, err)
+
+	var spec2 Spec
+	err = yaml.Unmarshal(dt, &spec2)
+	assert.NilError(t, err)
+
+	assert.Check(t, cmp.Equal(spec2.Name, "test"), spec2)
+	checkExt(t, spec2, "hello", "world")
+	checkExt(t, spec, "x-hello", "world")
+	checkExt(t, spec2, "foo", []string{"bar", "baz"})
+	checkExt(t, spec, "x-foo", []string{"bar", "baz"})
+	checkExt(t, spec, "capitalized", "world2")
+	checkExt(t, spec, "X-capitalized", "world2")
+}

spec.go

@@ -3,10 +3,14 @@ package dalec
 
 import (
 	"io/fs"
+	"strings"
 	"time"
 
+	"github.com/goccy/go-yaml"
+	"github.com/goccy/go-yaml/ast"
 	"github.com/moby/buildkit/client/llb"
 	"github.com/opencontainers/go-digest"
+	"github.com/pkg/errors"
 )
 
 // Spec is the specification for a package build.
@@ -94,6 +98,14 @@ type Spec struct {
 	// Each item in this list is run with a separate rootfs and cannot interact with other tests.
 	// Each [TestSpec] is run with a separate rootfs, asynchronously from other [TestSpec].
 	Tests []*TestSpec `yaml:"tests,omitempty" json:"tests,omitempty"`
+
+	// extRaw is the raw AST of the extension fields in the spec.
+	// This is used to extract the ext fields in [Spec.Ext]
+	extRaw *ast.File
+	// ext is all the ext fields in the spec.
+	// This gets used when marshalling the spec back to yaml.
+	// It is used to avoid having to re-parse the raw AST.
+	ext map[string]interface{}
 }
 
 // PatchSpec is used to apply a patch to a source with a given set of options.
@@ -427,3 +439,38 @@ func (s *SystemdConfiguration) EnabledUnits() map[string]SystemdUnitConfig {
 
 	return units
 }
+
+// Ext reads the extension field from the spec and unmarshals it into the target
+// value.
+func (s Spec) Ext(key string, target interface{}) error {
+	if !strings.HasPrefix(key, "x-") && !strings.HasPrefix(key, "X-") {
+		_, ok := s.ext["x-"+key]
+		if ok {
+			key = "x-" + key
+		} else {
+			_, ok = s.ext["X-"+key]
+			if !ok {
+				return errors.Errorf("extension field %q not found", key)
+			}
+			key = "X-" + key
+		}
+	}
+
+	p, err := yaml.PathString("$." + key)
+	if err != nil {
+		return err
+	}
+
+	node, err := p.FilterFile(s.extRaw)
+	if err != nil {
+		return errors.Wrap(err, "error filtering node")
+	}
+
+	dt := node.String()
+	err = yaml.Unmarshal([]byte(dt), target)
+	if err != nil {
+		return errors.Wrapf(err, "error unmarshalling extension field %q into target", key)
+	}
+
+	return nil
+}