datasize, blueprint: add new datasizes.Size type and use in blueprints

pkg/blueprint/disk_customizations.go

@@ -9,15 +9,37 @@ import (
 	"slices"
 	"strings"
 
+	"github.com/osbuild/images/pkg/datasizes"
 	"github.com/osbuild/images/pkg/pathpolicy"
 )
 
 type DiskCustomization struct {
 	// TODO: Add partition table type: gpt or dos
-	MinSize    uint64                   `json:"minsize,omitempty" toml:"minsize,omitempty"`
+	MinSize    uint64
+	Partitions []PartitionCustomization
+}
+
+type diskCustomizationMarshaler struct {
+	// TODO: Add partition table type: gpt or dos
+	MinSize    datasizes.Size           `json:"minsize,omitempty" toml:"minsize,omitempty"`
 	Partitions []PartitionCustomization `json:"partitions,omitempty" toml:"partitions,omitempty"`
 }
 
+func (dc *DiskCustomization) UnmarshalJSON(data []byte) error {
+	var dcm diskCustomizationMarshaler
+	if err := json.Unmarshal(data, &dcm); err != nil {
+		return err
+	}
+	dc.MinSize = dcm.MinSize.Uint64()
+	dc.Partitions = dcm.Partitions
+
+	return nil
+}
+
+func (dc *DiskCustomization) UnmarshalTOML(data any) error {
+	return unmarshalTOMLviaJSON(dc, data)
+}
+
 // PartitionCustomization defines a single partition on a disk. The Type
 // defines the kind of "payload" for the partition: plain, lvm, or btrfs.
 //   - plain: the payload will be a filesystem on a partition (e.g. xfs, ext4).

pkg/blueprint/disk_customizations_test.go

@@ -5,10 +5,12 @@ import (
 	"testing"
 
 	"github.com/BurntSushi/toml"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
 	"github.com/osbuild/images/pkg/blueprint"
 	"github.com/osbuild/images/pkg/datasizes"
 	"github.com/osbuild/images/pkg/pathpolicy"
-	"github.com/stretchr/testify/assert"
 )
 
 func TestPartitioningValidation(t *testing.T) {
@@ -1648,3 +1650,62 @@ func TestPartitionCustomizationUnmarshalTOML(t *testing.T) {
 		})
 	}
 }
+
+func TestDiskCustomizationUnmarshalJSON(t *testing.T) {
+	type testCase struct {
+		inputJSON string
+		inputTOML string
+		expected  *blueprint.DiskCustomization
+	}
+
+	testCases := map[string]testCase{
+		"nothing": {
+			inputJSON: "{}",
+			inputTOML: "",
+			expected: &blueprint.DiskCustomization{
+				MinSize: 0,
+			},
+		},
+		"minsize/int": {
+			inputJSON: `{
+				"minsize": 1234
+			}`,
+			inputTOML: "minsize = 1234",
+			expected: &blueprint.DiskCustomization{
+				MinSize: 1234,
+			},
+		},
+		"minsize/str": {
+			inputJSON: `{
+				"minsize": "1234"
+			}`,
+			inputTOML: `minsize = "1234"`,
+			expected: &blueprint.DiskCustomization{
+				MinSize: 1234,
+			},
+		},
+		"minsize/str-with-unit": {
+			inputJSON: `{
+				"minsize": "1 GiB"
+			}`,
+			inputTOML: `minsize = "1 GiB"`,
+			expected: &blueprint.DiskCustomization{
+				MinSize: 1 * datasizes.GiB,
+			},
+		},
+	}
+
+	for name := range testCases {
+		tc := testCases[name]
+		t.Run(name, func(t *testing.T) {
+			var dc blueprint.DiskCustomization
+
+			err := json.Unmarshal([]byte(tc.inputJSON), &dc)
+			require.NoError(t, err)
+			assert.Equal(t, tc.expected, &dc)
+			err = toml.Unmarshal([]byte(tc.inputTOML), &dc)
+			require.NoError(t, err)
+			assert.Equal(t, tc.expected, &dc)
+		})
+	}
+}

pkg/blueprint/filesystem_customizations.go

@@ -10,52 +10,33 @@ import (
 )
 
 type FilesystemCustomization struct {
-	Mountpoint string `json:"mountpoint,omitempty" toml:"mountpoint,omitempty"`
-	MinSize    uint64 `json:"minsize,omitempty" toml:"minsize,omitempty"`
+	Mountpoint string
+	MinSize    uint64
 }
 
-func (fsc *FilesystemCustomization) UnmarshalTOML(data interface{}) error {
-	d, ok := data.(map[string]interface{})
-	if !ok {
-		return fmt.Errorf("customizations.filesystem is not an object")
-	}
-
-	switch d["mountpoint"].(type) {
-	case string:
-		fsc.Mountpoint = d["mountpoint"].(string)
-	default:
-		return fmt.Errorf("TOML unmarshal: mountpoint must be string, got \"%v\" of type %T", d["mountpoint"], d["mountpoint"])
-	}
-	minSize, err := decodeSize(d["minsize"])
-	if err != nil {
-		return fmt.Errorf("TOML unmarshal: error decoding minsize value for mountpoint %q: %w", fsc.Mountpoint, err)
-	}
-	fsc.MinSize = minSize
-	return nil
+type filesystemCustomizationMarshaling struct {
+	Mountpoint string         `json:"mountpoint,omitempty" toml:"mountpoint,omitempty"`
+	MinSize    datasizes.Size `json:"minsize,omitempty" toml:"minsize,omitempty"`
 }
 
 func (fsc *FilesystemCustomization) UnmarshalJSON(data []byte) error {
-	var v interface{}
-	if err := json.Unmarshal(data, &v); err != nil {
+	var fc filesystemCustomizationMarshaling
+	if err := json.Unmarshal(data, &fc); err != nil {
+		if fc.Mountpoint != "" {
+			return fmt.Errorf("error decoding minsize value for mountpoint %q: %w", fc.Mountpoint, err)
+		}
 		return err
 	}
-	d, _ := v.(map[string]interface{})
+	fsc.Mountpoint = fc.Mountpoint
+	fsc.MinSize = fc.MinSize.Uint64()
 
-	switch d["mountpoint"].(type) {
-	case string:
-		fsc.Mountpoint = d["mountpoint"].(string)
-	default:
-		return fmt.Errorf("JSON unmarshal: mountpoint must be string, got \"%v\" of type %T", d["mountpoint"], d["mountpoint"])
-	}
-
-	minSize, err := decodeSize(d["minsize"])
-	if err != nil {
-		return fmt.Errorf("JSON unmarshal: error decoding minsize value for mountpoint %q: %w", fsc.Mountpoint, err)
-	}
-	fsc.MinSize = minSize
 	return nil
 }
 
+func (fsc *FilesystemCustomization) UnmarshalTOML(data any) error {
+	return unmarshalTOMLviaJSON(fsc, data)
+}
+
 // CheckMountpointsPolicy checks if the mountpoints are allowed by the policy
 func CheckMountpointsPolicy(mountpoints []FilesystemCustomization, mountpointAllowList *pathpolicy.PathPolicies) error {
 	var errs []error

pkg/blueprint/filesystem_customizations_test.go

@@ -47,19 +47,19 @@ func TestFilesystemCustomizationUnmarshalTOMLUnhappy(t *testing.T) {
 			name: "mountpoint not string",
 			input: `mountpoint = 42
 			minsize = 42`,
-			err: `toml: line 0: TOML unmarshal: mountpoint must be string, got "42" of type int64`,
+			err: `toml: line 0: error decoding TOML map[minsize:42 mountpoint:42]: json: cannot unmarshal number into Go struct field filesystemCustomizationMarshaling.mountpoint of type string`,
 		},
 		{
 			name: "minsize nor string nor int",
 			input: `mountpoint="/"
 			minsize = true`,
-			err: `toml: line 0: TOML unmarshal: error decoding minsize value for mountpoint "/": failed to convert value "true" to number`,
+			err: `toml: line 0: error decoding TOML map[minsize:true mountpoint:/]: error decoding size: failed to convert value "true" to number`,
 		},
 		{
 			name: "minsize not parseable",
 			input: `mountpoint="/"
 			minsize = "20 KG"`,
-			err: `toml: line 0: TOML unmarshal: error decoding minsize value for mountpoint "/": unknown data size units in string: 20 KG`,
+			err: `toml: line 0: error decoding TOML map[minsize:20 KG mountpoint:/]: error decoding size: unknown data size units in string: 20 KG`,
 		},
 	}
 
@@ -81,17 +81,17 @@ func TestFilesystemCustomizationUnmarshalJSONUnhappy(t *testing.T) {
 		{
 			name:  "mountpoint not string",
 			input: `{"mountpoint": 42, "minsize": 42}`,
-			err:   `JSON unmarshal: mountpoint must be string, got "42" of type float64`,
+			err:   `json: cannot unmarshal number into Go struct field filesystemCustomizationMarshaling.mountpoint of type string`,
 		},
 		{
 			name:  "minsize nor string nor int",
 			input: `{"mountpoint":"/", "minsize": true}`,
-			err:   `JSON unmarshal: error decoding minsize value for mountpoint "/": failed to convert value "true" to number`,
+			err:   `error decoding minsize value for mountpoint "/": error decoding size: failed to convert value "true" to number`,
 		},
 		{
 			name:  "minsize not parseable",
 			input: `{ "mountpoint": "/", "minsize": "20 KG"}`,
-			err:   `JSON unmarshal: error decoding minsize value for mountpoint "/": unknown data size units in string: 20 KG`,
+			err:   `error decoding minsize value for mountpoint "/": error decoding size: unknown data size units in string: 20 KG`,
 		},
 	}
 
@@ -115,7 +115,7 @@ func TestFilesystemCustomizationUnmarshalTOMLNotAnObject(t *testing.T) {
 			input: `
 [customizations]
 filesystem = ["hello"]`,
-			err: "toml: line 3 (last key \"customizations.filesystem\"): customizations.filesystem is not an object",
+			err: `toml: line 3 (last key "customizations.filesystem"): error decoding TOML hello: json: cannot unmarshal string into Go value of type blueprint.filesystemCustomizationMarshaling`,
 		},
 	}
 

pkg/blueprint/toml_json_bridge.go

@@ -0,0 +1,24 @@
+package blueprint
+
+import (
+	"encoding/json"
+	"fmt"
+)
+
+// XXX: move to interal/common ?
+func unmarshalTOMLviaJSON(u json.Unmarshaler, data any) error {
+	// This is the most efficient way to reuse code when unmarshaling
+	// structs in toml, it leaks json errors which is a bit sad but
+	// because the toml unmarshaler gives us not "[]byte" but an
+	// already pre-processed "any" we cannot just unmarshal into our
+	// "fooMarshaling" struct and reuse the result so we resort to
+	// this workaround (but toml will go away long term anyway).
+	dataJSON, err := json.Marshal(data)
+	if err != nil {
+		return fmt.Errorf("error unmarshaling TOML data %v: %w", data, err)
+	}
+	if err := u.UnmarshalJSON(dataJSON); err != nil {
+		return fmt.Errorf("error decoding TOML %v: %w", data, err)
+	}
+	return nil
+}

pkg/datasizes/size.go

@@ -0,0 +1,72 @@
+package datasizes
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+)
+
+// Size is a wrapper around uint64 with support for reading from string
+// yaml/toml, so {"size": 123}, {"size": "1234"}, {"size": "1 GiB"} are
+// all supported
+type Size uint64
+
+// Uint64 returns the size as uint64. This is a convenience functions,
+// it is strictly equivalent to uint64(Size(1))
+func (si Size) Uint64() uint64 {
+	return uint64(si)
+}
+
+func (si *Size) UnmarshalTOML(data interface{}) error {
+	i, err := decodeSize(data)
+	if err != nil {
+		return fmt.Errorf("error decoding TOML size: %w", err)
+	}
+	*si = Size(i)
+	return nil
+}
+
+func (si *Size) UnmarshalJSON(data []byte) error {
+	dec := json.NewDecoder(bytes.NewBuffer(data))
+	dec.UseNumber()
+
+	var v interface{}
+	if err := dec.Decode(&v); err != nil {
+		return err
+	}
+	i, err := decodeSize(v)
+	if err != nil {
+		// if only we could do better here and include e.g. the field
+		// name where this happend but encoding/json does not
+		// support this, c.f. https://github.com/golang/go/issues/58655
+		return fmt.Errorf("error decoding size: %w", err)
+	}
+	*si = Size(i)
+	return nil
+}
+
+// decodeSize takes an integer or string representing a data size (with a data
+// suffix) and returns the uint64 representation.
+func decodeSize(size any) (uint64, error) {
+	switch s := size.(type) {
+	case string:
+		return Parse(s)
+	case json.Number:
+		i, err := s.Int64()
+		if i < 0 {
+			return 0, fmt.Errorf("cannot be negative")
+		}
+		return uint64(i), err
+	case int64:
+		if s < 0 {
+			return 0, fmt.Errorf("cannot be negative")
+		}
+		return uint64(s), nil
+	case uint64:
+		return s, nil
+	case float64, float32:
+		return 0, fmt.Errorf("cannot be float")
+	default:
+		return 0, fmt.Errorf("failed to convert value \"%v\" to number", size)
+	}
+}