Merge #33221 #33222

33221: sql: always use typed expressions for shift on int r=mjibson a=mjibson

Previously the expressions 1 << 63 and 1::int << 63 would return positive
and negative results, respectively. This was because constant folding
in the first case is not subject to 2s complent that causes negative
ints in the second case.

Change shift on int to not fold. This prevents the surprising behavior
described above in addition to other bugs where shifting by a large
amount could allocate huge amounts of memory and crash the process.

Fixes #32682

Release note (sql change): Shift operators on integers now always operate
on explicit types instead of possibly an untyped constant. They also
now return an error if the right-side operand is out of range.

33222: pgwire: test \u0001 in JSON pgwire encoding r=mjibson a=mjibson

Verify that we handle "\u0001" JSON strings identically to postgres. The
printing of \u0001 instead of some other character is correct, and the
same as what postgres does. (We do not test \u0000 here because postgres
rejects it as a valid string, even though we accept it.)

Closes #33165

Release note: None

Co-authored-by: Matt Jibson <matt.jibson@gmail.com>

pkg/cmd/generate-binary/main.go

@@ -355,6 +355,7 @@ var inputs = map[string][]string{
 		`false`,
 		`null`,
 		`[[[[true, false, null]]]]`,
+		`["\u0001", "\u0041", "\u26a3", "\ud83e\udd37"]`,
 	},
 
 	"'%s'::uuid[]": {

pkg/sql/logictest/testdata/logic_test/shift

@@ -0,0 +1,49 @@
+# LogicTest: local local-opt local-parallel-stmts fakedist fakedist-opt fakedist-metadata
+
+# Check non-constant eval
+
+statement ok
+CREATE TABLE t AS SELECT 1 AS i
+
+statement error shift argument out of range
+SELECT i << 64 FROM t
+
+statement error shift argument out of range
+SELECT i >> 64 FROM t
+
+statement error shift argument out of range
+SELECT i << -1 FROM t
+
+statement error shift argument out of range
+SELECT i >> -1 FROM t
+
+query II
+SELECT i << 63 >> 63, i << 62 >> 62 FROM t
+----
+-1 1
+
+# Check constant folding
+
+statement error shift argument out of range
+SELECT 1 << 64
+
+statement error shift argument out of range
+SELECT 1 >> 64
+
+statement error shift argument out of range
+SELECT 1 << -1
+
+statement error shift argument out of range
+SELECT 1 >> -1
+
+query II
+SELECT 1 << 63 >> 63, 1 << 62 >> 62
+----
+-1 1
+
+# Ensure that shift returns the same result as an int or a constant
+
+query II
+SELECT 1 << 63 >> 63, 1::INT << 63 >> 63
+----
+-1 -1

pkg/sql/pgwire/testdata/encodings.json

@@ -989,6 +989,12 @@
 		"TextAsBinary": [91, 91, 91, 91, 116, 114, 117, 101, 44, 32, 102, 97, 108, 115, 101, 44, 32, 110, 117, 108, 108, 93, 93, 93, 93],
 		"Binary": [1, 91, 91, 91, 91, 116, 114, 117, 101, 44, 32, 102, 97, 108, 115, 101, 44, 32, 110, 117, 108, 108, 93, 93, 93, 93]
 	},
+	{
+		"SQL": "'[\"\\u0001\", \"\\u0041\", \"\\u26a3\", \"\\ud83e\\udd37\"]'::jsonb",
+		"Text": "[\"\\u0001\", \"A\", \"⚣\", \"🤷\"]",
+		"TextAsBinary": [91, 34, 92, 117, 48, 48, 48, 49, 34, 44, 32, 34, 65, 34, 44, 32, 34, 226, 154, 163, 34, 44, 32, 34, 240, 159, 164, 183, 34, 93],
+		"Binary": [1, 91, 34, 92, 117, 48, 48, 48, 49, 34, 44, 32, 34, 65, 34, 44, 32, 34, 226, 154, 163, 34, 44, 32, 34, 240, 159, 164, 183, 34, 93]
+	},
 	{
 		"SQL": "'00:00:00'::time",
 		"Text": "00:00:00",

pkg/sql/sem/tree/constant.go

@@ -530,10 +530,6 @@ var binaryOpToTokenIntOnly = map[BinaryOperator]token.Token{
 	Bitor:    token.OR,
 	Bitxor:   token.XOR,
 }
-var binaryShiftOpToToken = map[BinaryOperator]token.Token{
-	LShift: token.SHL,
-	RShift: token.SHR,
-}
 var comparisonOpToToken = map[ComparisonOperator]token.Token{
 	EQ: token.EQL,
 	NE: token.NEQ,
@@ -594,13 +590,10 @@ func (constantFolderVisitor) VisitPost(expr Expr) (retExpr Expr) {
 						}
 					}
 				}
-				if token, ok := binaryShiftOpToToken[t.Operator]; ok {
-					if lInt, ok := l.asConstantInt(); ok {
-						if rInt64, err := r.AsInt64(); err == nil && rInt64 >= 0 {
-							return &NumVal{Value: constant.Shift(lInt, token, uint(rInt64))}
-						}
-					}
-				}
+				// Explicitly ignore shift operators so the expression is evaluated as a
+				// non-const. This is because 1 << 63 as a 64-bit int (which is a negative
+				// number due to 2s complement) is different than 1 << 63 as constant,
+				// which is positive.
 			}
 		case *StrVal:
 			if r, ok := t.Right.(*StrVal); ok {

pkg/sql/sem/tree/constant_test.go

@@ -445,11 +445,11 @@ func TestFoldNumericConstants(t *testing.T) {
 		{`2 ^ 3`, `2 ^ 3`},         // Constant folding won't fold power.
 		{`1.3 ^ 3.9`, `1.3 ^ 3.9`},
 		// Shift ops (int only).
-		{`1 << 2`, `4`},
+		{`1 << 2`, `1 << 2`},
 		{`1 << -2`, `1 << -2`},                                                     // Should be caught during evaluation.
 		{`1 << 9999999999999999999999999999`, `1 << 9999999999999999999999999999`}, // Will be caught during type checking.
 		{`1.2 << 2.4`, `1.2 << 2.4`},                                               // Will be caught during type checking.
-		{`4 >> 2`, `1`},
+		{`4 >> 2`, `4 >> 2`},
 		{`4.1 >> 2.9`, `4.1 >> 2.9`}, // Will be caught during type checking.
 		// Comparison ops.
 		{`4 = 2`, `false`},
@@ -487,7 +487,7 @@ func TestFoldNumericConstants(t *testing.T) {
 		{`(((4)))`, `4`},
 		{`(((9 / 3) * (1 / 3)))`, `1`},
 		{`(((9 / 3) % (1 / 3)))`, `((3 % 0.333333))`},
-		{`(1.0) << ((2) + 3 / (1/9))`, `536870912`},
+		{`(1.0) << ((2) + 3 / (1/9))`, `1.0 << 29`},
 		// With non-constants.
 		{`a + 5 * b`, `a + (5 * b)`},
 		{`a + 5 + b + 7`, `((a + 5) + b) + 7`},

pkg/sql/sem/tree/eval.go

@@ -29,6 +29,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/internal/client"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/server/telemetry"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql/coltypes"
 	"github.com/cockroachdb/cockroach/pkg/sql/lex"
@@ -1324,7 +1325,12 @@ var BinOps = map[BinaryOperator]binOpOverload{
 			RightType:  types.Int,
 			ReturnType: types.Int,
 			Fn: func(_ *EvalContext, left Datum, right Datum) (Datum, error) {
-				return NewDInt(MustBeDInt(left) << uint(MustBeDInt(right))), nil
+				rval := MustBeDInt(right)
+				if rval < 0 || rval >= 64 {
+					telemetry.Count("sql.large_lshift_argument")
+					return nil, pgerror.NewErrorf(pgerror.CodeInvalidParameterValueError, "shift argument out of range")
+				}
+				return NewDInt(MustBeDInt(left) << uint(rval)), nil
 			},
 		},
 		&BinOp{
@@ -1357,7 +1363,12 @@ var BinOps = map[BinaryOperator]binOpOverload{
 			RightType:  types.Int,
 			ReturnType: types.Int,
 			Fn: func(_ *EvalContext, left Datum, right Datum) (Datum, error) {
-				return NewDInt(MustBeDInt(left) >> uint(MustBeDInt(right))), nil
+				rval := MustBeDInt(right)
+				if rval < 0 || rval >= 64 {
+					telemetry.Count("sql.large_rshift_argument")
+					return nil, pgerror.NewErrorf(pgerror.CodeInvalidParameterValueError, "shift argument out of range")
+				}
+				return NewDInt(MustBeDInt(left) >> uint(rval)), nil
 			},
 		},
 		&BinOp{