fix: properly parse matcher expressions

Fixes: siderolabs#3572 Signed-off-by: Artem Chernyshev <artem.0xD2@gmail.com>
smira · May 4, 2021 · 3823909 · 3823909
1 parent e54b6b7
commit 3823909
Show file tree

Hide file tree

Showing 3 changed files with 146 additions and 32 deletions.
diff --git a/pkg/machinery/config/types/v1alpha1/v1alpha1_marshal_test.go b/pkg/machinery/config/types/v1alpha1/v1alpha1_marshal_test.go
@@ -5,10 +5,13 @@
 package v1alpha1_test
 
 import (
+	"fmt"
 	"testing"
 
+	humanize "github.com/dustin/go-humanize"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"github.com/talos-systems/go-blockdevice/blockdevice/util/disk"
 	"gopkg.in/yaml.v3"
 
 	"github.com/talos-systems/talos/pkg/machinery/config/types/v1alpha1"
@@ -35,3 +38,112 @@ func TestBase64Bytes(t *testing.T) {
 
 	assert.Equal(t, input.CA, decoded.CA)
 }
+
+func TestDiskSizeMatcherUnmarshal(t *testing.T) {
+	obj := struct {
+		M *v1alpha1.InstallDiskSizeMatcher `yaml:"m"`
+	}{
+		M: &v1alpha1.InstallDiskSizeMatcher{},
+	}
+
+	for _, test := range []struct {
+		condition string
+		size      string
+		match     bool
+		err       bool
+	}{
+		{
+			condition: "<= 256GB",
+			size:      "200GB",
+			match:     true,
+		},
+		{
+			condition: ">= 256GB",
+			size:      "200GB",
+			match:     false,
+		},
+		{
+			condition: "<256GB",
+			size:      "256GB",
+			match:     false,
+		},
+		{
+			condition: ">256GB",
+			size:      "256GB",
+			match:     false,
+		},
+		{
+			condition: ">256GB",
+			size:      "257GB",
+			match:     true,
+		},
+		{
+			condition: "==256GB",
+			size:      "256GB",
+			match:     true,
+		},
+		{
+			condition: "==256GB",
+			size:      "257GB",
+			match:     false,
+		},
+		{
+			condition: "==   256GB",
+			size:      "256GB",
+			match:     true,
+		},
+		{
+			condition: "256GB",
+			size:      "256GB",
+			match:     true,
+		},
+		{
+			condition: "   256GB",
+			size:      "256GB",
+			match:     true,
+		},
+		{
+			condition: "9a256GB",
+			err:       true,
+		},
+		{
+			condition: "--256GB",
+			err:       true,
+		},
+		{
+			condition: "<<256GB",
+			err:       true,
+		},
+		{
+			condition: ">1",
+			size:      "1GB",
+			match:     true,
+		},
+		{
+			condition: "< 1",
+			size:      "1GB",
+			match:     false,
+		},
+	} {
+		var (
+			size uint64
+			err  error
+		)
+
+		if test.size != "" {
+			size, err = humanize.ParseBytes(test.size)
+			require.NoError(t, err)
+		}
+
+		err = yaml.Unmarshal([]byte(fmt.Sprintf("m: '%s'\n", test.condition)), &obj)
+		if test.err {
+			require.Error(t, err)
+		} else {
+			require.NoError(t, err)
+		}
+
+		if test.size != "" {
+			require.Equal(t, obj.M.Matcher(&disk.Disk{Size: size}), test.match, test.size)
+		}
+	}
+}
diff --git a/pkg/machinery/config/types/v1alpha1/v1alpha1_provider.go b/pkg/machinery/config/types/v1alpha1/v1alpha1_provider.go
@@ -858,7 +858,7 @@ func (i *InstallConfig) DiskMatchers() []disk.Matcher {
 
 		matchers := []disk.Matcher{}
 		if selector.Size != nil {
-			matchers = append(matchers, selector.Size.matcher)
+			matchers = append(matchers, selector.Size.Matcher)
 		}
 
 		if selector.UUID != "" {

diff --git a/pkg/machinery/config/types/v1alpha1/v1alpha1_types.go b/pkg/machinery/config/types/v1alpha1/v1alpha1_types.go
@@ -846,7 +846,7 @@ type InstallConfig struct {
 // InstallDiskSizeMatcher disk size condition parser.
 type InstallDiskSizeMatcher struct {
 	// docgen:nodoc
-	matcher disk.Matcher
+	Matcher disk.Matcher
 	// docgen:nodoc
 	condition string
 }
@@ -856,54 +856,56 @@ func (m *InstallDiskSizeMatcher) MarshalYAML() (interface{}, error) {
 	return m.condition, nil
 }
 
-// UnmarshalYAML is a custom unmarshaller for `Endpoint`.
+// UnmarshalYAML is a custom unmarshaller for `InstallDiskSizeMatcher`.
 func (m *InstallDiskSizeMatcher) UnmarshalYAML(unmarshal func(interface{}) error) error {
 	if err := unmarshal(&m.condition); err != nil {
 		return err
 	}
 
-	re := regexp.MustCompile(`(>=|<=|>|<|==)?\b*(.+)$`)
+	m.condition = strings.TrimSpace(m.condition)
+
+	re := regexp.MustCompile(`(>=|<=|>|<|==)?\b*(.*)$`)
 
 	parts := re.FindStringSubmatch(m.condition)
-	if len(parts) == 0 {
+	if len(parts) < 2 {
 		return fmt.Errorf("failed to parse the condition: expected [>=|<=|>|<|==]<size>[units], got %s", m.condition)
 	}
 
 	var compare func(*disk.Disk, uint64) bool
 
-	if len(parts) > 1 {
-		switch parts[0] {
-		case ">=":
-			compare = func(d *disk.Disk, size uint64) bool {
-				return d.Size >= size
-			}
-		case "<=":
-			compare = func(d *disk.Disk, size uint64) bool {
-				return d.Size <= size
-			}
-		case ">":
-			compare = func(d *disk.Disk, size uint64) bool {
-				return d.Size > size
-			}
-		case "<":
-			compare = func(d *disk.Disk, size uint64) bool {
-				return d.Size < size
-			}
-		case "==":
-			compare = func(d *disk.Disk, size uint64) bool {
-				return d.Size == size
-			}
-		default:
-			return fmt.Errorf("unknown binary operator %s", parts[0])
+	switch parts[1] {
+	case ">=":
+		compare = func(d *disk.Disk, size uint64) bool {
+			return d.Size >= size
+		}
+	case "<=":
+		compare = func(d *disk.Disk, size uint64) bool {
+			return d.Size <= size
+		}
+	case ">":
+		compare = func(d *disk.Disk, size uint64) bool {
+			return d.Size > size
 		}
+	case "<":
+		compare = func(d *disk.Disk, size uint64) bool {
+			return d.Size < size
+		}
+	case "":
+		fallthrough
+	case "==":
+		compare = func(d *disk.Disk, size uint64) bool {
+			return d.Size == size
+		}
+	default:
+		return fmt.Errorf("unknown binary operator %s", parts[1])
 	}
 
-	size, err := humanize.ParseBytes(parts[1])
+	size, err := humanize.ParseBytes(strings.TrimSpace(parts[2]))
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to parse disk size %s: %s", parts[2], err)
 	}
 
-	m.matcher = func(d *disk.Disk) bool {
+	m.Matcher = func(d *disk.Disk) bool {
 		return compare(d, size)
 	}