provenance: avoid intermediate wrapper for custom fields

While the custom fields used embedded Go struct field,
this does not work for map types and needs custom JSON
marshaller to make sure custom fields appear without wrapper.

Signed-off-by: Tonis Tiigi <[email protected]>

Author: tonistiigi

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))
+}