fix: introduce gRPC input validation

Validate that both label & ID queries have sane limits on inputs.

For label query - ensure that values are populated, for the ID query
ensure a maximum budget for the regexp to be compiled. In real life ID
queries should be simple (and probably regexp is an overkill here), so
limit the regexp maximum complexity.

Signed-off-by: Andrey Smirnov <andrey.smirnov@siderolabs.com>

Author: smira

pkg/state/protobuf/server/helpers.go

@@ -6,6 +6,7 @@ package server
 
 import (
 	"regexp"
+	"regexp/syntax"
 
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -25,6 +26,19 @@ func ConvertLabelQuery(terms []*v1alpha1.LabelTerm) ([]resource.LabelQueryOption
 			opts = append(opts, resource.NotMatches)
 		}
 
+		// Value is a repeated field in the protobuf definition, but the operators below
+		// require a single value, so reject the request instead of panicking on an empty one.
+		switch term.Op { //nolint:exhaustive
+		case v1alpha1.LabelTerm_EQUAL,
+			v1alpha1.LabelTerm_LT,
+			v1alpha1.LabelTerm_LTE,
+			v1alpha1.LabelTerm_LT_NUMERIC,
+			v1alpha1.LabelTerm_LTE_NUMERIC:
+			if len(term.Value) == 0 {
+				return nil, status.Errorf(codes.InvalidArgument, "missing value for label query operator: %v", term.Op)
+			}
+		}
+
 		switch term.Op {
 		case v1alpha1.LabelTerm_EQUAL:
 			labelOpts = append(labelOpts, resource.LabelEqual(term.Key, term.Value[0], opts...))
@@ -50,16 +64,95 @@ func ConvertLabelQuery(terms []*v1alpha1.LabelTerm) ([]resource.LabelQueryOption
 	return labelOpts, nil
 }
 
+const (
+	// maxRegexpLength is the maximum length of the ID query regexp.
+	//
+	// Parsing cost is roughly linear in the length of the pattern, so this bounds the
+	// work done before the program size below can be estimated.
+	maxRegexpLength = 4096
+
+	// maxRegexpProgramSize is the maximum size (in instructions) of the compiled ID query regexp.
+	//
+	// Repetition expands a pattern far beyond its length: the parser bounds the expansion
+	// only by its own ~3.3M instruction limit, so the length limit above is not by itself a
+	// useful bound on the compiled program. Reaching that ceiling costs ~700MiB and a few
+	// hundred milliseconds per query.
+	maxRegexpProgramSize = 4096
+)
+
 // ConvertIDQuery converts protobuf representation of IDQuery to state representation.
 func ConvertIDQuery(input *v1alpha1.IDQuery) ([]resource.IDQueryOption, error) {
 	if input == nil || input.Regexp == "" {
 		return nil, nil
 	}
 
+	if len(input.Regexp) > maxRegexpLength {
+		return nil, status.Errorf(codes.InvalidArgument, "regexp is too long: %d > %d", len(input.Regexp), maxRegexpLength)
+	}
+
+	parsed, err := syntax.Parse(input.Regexp, syntax.Perl)
+	if err != nil {
+		return nil, status.Errorf(codes.InvalidArgument, "failed to parse regexp: %v", err)
+	}
+
+	if size := regexpProgramSize(parsed); size > maxRegexpProgramSize {
+		return nil, status.Errorf(codes.InvalidArgument, "regexp is too complex: %d > %d", size, maxRegexpProgramSize)
+	}
+
 	re, err := regexp.Compile(input.Regexp)
 	if err != nil {
 		return nil, status.Errorf(codes.InvalidArgument, "failed to compile regexp: %v", err)
 	}
 
 	return []resource.IDQueryOption{resource.IDRegexpMatch(re)}, nil
 }
+
+// regexpProgramSize estimates the number of instructions the parsed regexp compiles to.
+//
+// This is a port of (*parser).calcSize from regexp/syntax, which the parser itself uses to
+// enforce its (much higher) built-in limit. It is accurate to within the couple of extra
+// instructions the compiler emits for the program as a whole, and unlike compiling the
+// regexp it doesn't allocate the program. Recursion is bounded by the nesting depth limit
+// enforced by the parser.
+func regexpProgramSize(re *syntax.Regexp) int64 {
+	var size int64
+
+	switch re.Op { //nolint:exhaustive
+	case syntax.OpLiteral:
+		size = int64(len(re.Rune))
+	case syntax.OpCapture, syntax.OpStar:
+		// star can be 1+ or 2+; assume 2 pessimistically
+		size = 2 + regexpProgramSize(re.Sub[0])
+	case syntax.OpPlus, syntax.OpQuest:
+		size = 1 + regexpProgramSize(re.Sub[0])
+	case syntax.OpConcat:
+		for _, sub := range re.Sub {
+			size += regexpProgramSize(sub)
+		}
+	case syntax.OpAlternate:
+		for _, sub := range re.Sub {
+			size += regexpProgramSize(sub)
+		}
+
+		if len(re.Sub) > 1 {
+			size += int64(len(re.Sub)) - 1
+		}
+	case syntax.OpRepeat:
+		sub := regexpProgramSize(re.Sub[0])
+
+		if re.Max == -1 {
+			if re.Min == 0 {
+				size = 2 + sub // x*
+			} else {
+				size = 1 + int64(re.Min)*sub // xxx+
+			}
+
+			break
+		}
+
+		// x{2,5} = xx(x(x(x)?)?)?
+		size = int64(re.Max)*sub + int64(re.Max-re.Min)
+	}
+
+	return max(1, size)
+}

pkg/state/protobuf/server/helpers_test.go

@@ -0,0 +1,164 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+package server_test
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	"github.com/cosi-project/runtime/api/v1alpha1"
+	"github.com/cosi-project/runtime/pkg/state/protobuf/server"
+)
+
+func TestConvertIDQuery(t *testing.T) {
+	t.Parallel()
+
+	for _, test := range []struct {
+		name        string
+		regexp      string
+		expectedErr string
+		expectedNil bool
+	}{
+		{
+			name:        "empty",
+			regexp:      "",
+			expectedNil: true,
+		},
+		{
+			name:   "simple",
+			regexp: "^foo-[0-9]+$",
+		},
+		{
+			name:        "invalid",
+			regexp:      "^foo-[0-9+$",
+			expectedErr: "rpc error: code = InvalidArgument desc = failed to parse regexp: error parsing regexp: missing closing ]: `[0-9+$`",
+		},
+		{
+			name:        "too long",
+			regexp:      strings.Repeat("a", 4097),
+			expectedErr: "rpc error: code = InvalidArgument desc = regexp is too long: 4097 > 4096",
+		},
+		{
+			// short pattern, but expands to ~1M instructions and ~200MiB when compiled
+			name:        "too complex",
+			regexp:      "(" + strings.Repeat("a", 1024) + "){1000}",
+			expectedErr: "rpc error: code = InvalidArgument desc = regexp is too complex: 1026000 > 4096",
+		},
+		{
+			name:        "too complex, nested repeats",
+			regexp:      "((abcdefghij){10}){100}",
+			expectedErr: "rpc error: code = InvalidArgument desc = regexp is too complex: 12200 > 4096",
+		},
+		{
+			name:   "just under the complexity limit",
+			regexp: "(" + strings.Repeat("a", 7) + "){450}",
+		},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			opts, err := server.ConvertIDQuery(&v1alpha1.IDQuery{Regexp: test.regexp})
+
+			if test.expectedErr != "" {
+				require.Error(t, err)
+				assert.Equal(t, test.expectedErr, err.Error())
+				assert.Nil(t, opts)
+
+				return
+			}
+
+			require.NoError(t, err)
+
+			if test.expectedNil {
+				assert.Nil(t, opts)
+			} else {
+				assert.Len(t, opts, 1)
+			}
+		})
+	}
+}
+
+func TestConvertIDQueryNil(t *testing.T) {
+	t.Parallel()
+
+	opts, err := server.ConvertIDQuery(nil)
+	require.NoError(t, err)
+	assert.Nil(t, opts)
+}
+
+func TestConvertLabelQueryMissingValue(t *testing.T) {
+	t.Parallel()
+
+	for _, op := range []v1alpha1.LabelTerm_Operation{
+		v1alpha1.LabelTerm_EQUAL,
+		v1alpha1.LabelTerm_LT,
+		v1alpha1.LabelTerm_LTE,
+		v1alpha1.LabelTerm_LT_NUMERIC,
+		v1alpha1.LabelTerm_LTE_NUMERIC,
+	} {
+		t.Run(op.String(), func(t *testing.T) {
+			t.Parallel()
+
+			opts, err := server.ConvertLabelQuery([]*v1alpha1.LabelTerm{{Op: op, Key: "key"}})
+			require.Error(t, err)
+			assert.Equal(t, "rpc error: code = InvalidArgument desc = missing value for label query operator: "+op.String(), err.Error())
+			assert.Nil(t, opts)
+		})
+	}
+}
+
+func TestConvertLabelQuery(t *testing.T) {
+	t.Parallel()
+
+	// operators which don't require a value
+	for _, op := range []v1alpha1.LabelTerm_Operation{
+		v1alpha1.LabelTerm_EXISTS,
+		v1alpha1.LabelTerm_NOT_EXISTS,
+		v1alpha1.LabelTerm_IN,
+	} {
+		t.Run(op.String(), func(t *testing.T) {
+			t.Parallel()
+
+			opts, err := server.ConvertLabelQuery([]*v1alpha1.LabelTerm{{Op: op, Key: "key"}})
+			require.NoError(t, err)
+			assert.Len(t, opts, 1)
+		})
+	}
+
+	t.Run("with values", func(t *testing.T) {
+		t.Parallel()
+
+		opts, err := server.ConvertLabelQuery([]*v1alpha1.LabelTerm{
+			{Op: v1alpha1.LabelTerm_EQUAL, Key: "key", Value: []string{"value"}},
+			{Op: v1alpha1.LabelTerm_LTE_NUMERIC, Key: "key", Value: []string{"5"}, Invert: true},
+		})
+		require.NoError(t, err)
+		assert.Len(t, opts, 2)
+	})
+
+	t.Run("missing values", func(t *testing.T) {
+		t.Parallel()
+
+		_, err := server.ConvertLabelQuery([]*v1alpha1.LabelTerm{
+			{Op: v1alpha1.LabelTerm_EQUAL, Key: "key"},
+		})
+		require.Error(t, err)
+		assert.Equal(t, codes.InvalidArgument, status.Code(err))
+	})
+
+	t.Run("unsupported operator", func(t *testing.T) {
+		t.Parallel()
+
+		opts, err := server.ConvertLabelQuery([]*v1alpha1.LabelTerm{{Op: v1alpha1.LabelTerm_Operation(100), Key: "key"}})
+		require.Error(t, err)
+		assert.Equal(t, "rpc error: code = Unimplemented desc = unsupported label query operator: 100", err.Error())
+		assert.Nil(t, opts)
+	})
+}