provenance: avoid intermediate wrapper for custom fields

solver/llbsolver/provenance/types/types.go

@@ -1,6 +1,8 @@
 package types
 
 import (
+	"encoding/json"
+	"maps"
 	"slices"
 
 	slsa "github.com/in-toto/in-toto-golang/in_toto/slsa_provenance/common"
@@ -311,3 +313,95 @@ func (p *ProvenancePredicateSLSA02) ConvertToSLSA1() *ProvenancePredicateSLSA1 {
 		RunDetails:      runDetails,
 	}
 }
+
+// MarshalJSON flattens ProvenanceCustomEnv into top level.
+func (p ProvenanceInternalParametersSLSA1) MarshalJSON() ([]byte, error) {
+	type Alias ProvenanceInternalParametersSLSA1
+	base, err := json.Marshal(Alias(p))
+	if err != nil {
+		return nil, err
+	}
+	var m map[string]any
+	if err := json.Unmarshal(base, &m); err != nil {
+		return nil, err
+	}
+	maps.Copy(m, p.ProvenanceCustomEnv)
+	delete(m, "ProvenanceCustomEnv")
+	return json.Marshal(m)
+}
+
+// UnmarshalJSON fills both struct fields and flattened custom env.
+func (p *ProvenanceInternalParametersSLSA1) UnmarshalJSON(data []byte) error {
+	var m map[string]any
+	if err := json.Unmarshal(data, &m); err != nil {
+		return err
+	}
+
+	type Alias ProvenanceInternalParametersSLSA1
+	var a Alias
+	if err := json.Unmarshal(data, &a); err != nil {
+		return err
+	}
+
+	// Unmarshal known struct again to identify its keys
+	structBytes, err := json.Marshal(a)
+	if err != nil {
+		return err
+	}
+	var known map[string]any
+	if err := json.Unmarshal(structBytes, &known); err != nil {
+		return err
+	}
+
+	for k := range known {
+		delete(m, k)
+	}
+
+	*p = ProvenanceInternalParametersSLSA1(a)
+	p.ProvenanceCustomEnv = m
+	return nil
+}
+
+func (p Environment) MarshalJSON() ([]byte, error) {
+	type Alias Environment
+	base, err := json.Marshal(Alias(p))
+	if err != nil {
+		return nil, err
+	}
+	var m map[string]any
+	if err := json.Unmarshal(base, &m); err != nil {
+		return nil, err
+	}
+	maps.Copy(m, p.ProvenanceCustomEnv)
+	delete(m, "ProvenanceCustomEnv")
+	return json.Marshal(m)
+}
+
+func (p *Environment) UnmarshalJSON(data []byte) error {
+	var m map[string]any
+	if err := json.Unmarshal(data, &m); err != nil {
+		return err
+	}
+
+	type Alias Environment
+	var a Alias
+	if err := json.Unmarshal(data, &a); err != nil {
+		return err
+	}
+	// Unmarshal known struct again to identify its keys
+	structBytes, err := json.Marshal(a)
+	if err != nil {
+		return err
+	}
+	var known map[string]any
+	if err := json.Unmarshal(structBytes, &known); err != nil {
+		return err
+	}
+
+	for k := range known {
+		delete(m, k)
+	}
+	*p = Environment(a)
+	p.ProvenanceCustomEnv = m
+	return nil
+}

solver/llbsolver/provenance/types/types_test.go

@@ -0,0 +1,73 @@
+package types
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestMarsalBuildDefinitionSLSA1(t *testing.T) {
+	inp := `{
+"buildType": "btype1",
+"externalParameters": {
+	"configSource": {},
+	"request": {}
+},
+"internalParameters": {
+	"builderPlatform": "linux/amd64",
+		"foo": "bar",
+		"abc": 123,
+		"def": {"one": 1}
+	}
+}`
+
+	var def ProvenanceBuildDefinitionSLSA1
+	err := json.Unmarshal([]byte(inp), &def)
+	require.NoError(t, err)
+
+	require.Equal(t, "btype1", def.BuildType)
+	require.Equal(t, "linux/amd64", def.InternalParameters.BuilderPlatform)
+	require.Equal(t, "bar", def.InternalParameters.ProvenanceCustomEnv["foo"])
+	require.InEpsilon(t, float64(123), def.InternalParameters.ProvenanceCustomEnv["abc"], 0.001)
+	require.Equal(t, map[string]any{"one": float64(1)}, def.InternalParameters.ProvenanceCustomEnv["def"])
+
+	out, err := json.Marshal(def)
+	require.NoError(t, err)
+
+	require.JSONEq(t, inp, string(out))
+}
+
+func TestMarshalInvocation(t *testing.T) {
+	inp := `{
+	"configSource": {
+		"uri": "git+https://github.com/example/repo.git"
+	},
+	"parameters": {
+		"frontend": "dockerfile.v0"
+	},
+	"environment": {
+		"platform": "linux/amd64",
+		"buildkit": "v0.10.3",
+		"custom": {
+			"foo": "bar"
+		},
+		"bar": [1,2,3]
+	}
+}`
+
+	var inv ProvenanceInvocationSLSA02
+	err := json.Unmarshal([]byte(inp), &inv)
+	require.NoError(t, err)
+
+	require.Equal(t, "git+https://github.com/example/repo.git", inv.ConfigSource.URI)
+	require.Equal(t, "dockerfile.v0", inv.Parameters.Frontend)
+	require.Equal(t, "linux/amd64", inv.Environment.Platform)
+	require.Equal(t, "v0.10.3", inv.Environment.ProvenanceCustomEnv["buildkit"])
+	require.Equal(t, "bar", inv.Environment.ProvenanceCustomEnv["custom"].(map[string]any)["foo"])
+	require.Equal(t, []any{float64(1), float64(2), float64(3)}, inv.Environment.ProvenanceCustomEnv["bar"])
+	out, err := json.Marshal(inv)
+	require.NoError(t, err)
+
+	require.JSONEq(t, inp, string(out))
+}