Merge cockroachdb#108039 cockroachdb#108296

108039: sqlproxyccl: fix non-auth error causing throttle r=darinpp a=darinpp

Currently an error that occurs after a successful authentication but before the SQL server is ready to serve queries is considered an authentication failure and causes a throttle. In the cases where the connections are limited (due to usage limits) this causes blocks and prevents the clients to connect. This PR fixes the issue by not considering errors that came after the authentication was successful as conditions to throttle.

Fixes: https://cockroachlabs.atlassian.net/browse/CC-25314

Release note: None

108296: go.mod: bump Pebble to 40d3f411e45b r=RaduBerinde a=RaduBerinde

40d3f411 pebble: expose the secondary cache size as an option
077c3e8b ci: use any 1.20 go version
ff9402f0 Makefile: add testcoverage target
7e548be8 db: log flushed memtables in-memory size
83c9361c pebble: Collect Pebble Key Statistics
fe0ea048 manifest: Added Virtual SSTable Metadata Bounds Checking
42d81d23 internal/rangekey: remove ForeignSSTTransformer
9c487492 *: simplify pointCollapsingIter, use vanilla iters for foreign SSTs
5e7f8852 db: introduce Download API

Release note: None
Epic: none

Co-authored-by: Darin Peshev <[email protected]>
Co-authored-by: Radu Berinde <[email protected]>

Author: 3 people

DEPS.bzl

@@ -1595,10 +1595,10 @@ def go_deps():
         patches = [
             "@com_github_cockroachdb_cockroach//build/patches:com_github_cockroachdb_pebble.patch",
         ],
-        sha256 = "7ea06d5b0207ed3f20b5962bde7af18e6744fe089a98c51d66ac812252746342",
-        strip_prefix = "github.com/cockroachdb/[email protected]20230728153158-ce43e6535942",
+        sha256 = "f0319e618ed024b7d708e2c1f8cf0d4b9b7e4112943288ba6036e893a7f8c151",
+        strip_prefix = "github.com/cockroachdb/[email protected]20230807145728-40d3f411e45b",
         urls = [
-            "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/pebble/com_github_cockroachdb_pebble-v0.0.0-20230728153158-ce43e6535942.zip",
+            "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/pebble/com_github_cockroachdb_pebble-v0.0.0-20230807145728-40d3f411e45b.zip",
         ],
     )
     go_repository(

build/bazelutil/distdir_files.bzl

@@ -320,7 +320,7 @@ DISTDIR_FILES = {
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/go-test-teamcity/com_github_cockroachdb_go_test_teamcity-v0.0.0-20191211140407-cff980ad0a55.zip": "bac30148e525b79d004da84d16453ddd2d5cd20528e9187f1d7dac708335674b",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/gostdlib/com_github_cockroachdb_gostdlib-v1.19.0.zip": "c4d516bcfe8c07b6fc09b8a9a07a95065b36c2855627cb3514e40c98f872b69e",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/logtags/com_github_cockroachdb_logtags-v0.0.0-20230118201751-21c54148d20b.zip": "ca7776f47e5fecb4c495490a679036bfc29d95bd7625290cfdb9abb0baf97476",
-    "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/pebble/com_github_cockroachdb_pebble-v0.0.0-20230728153158-ce43e6535942.zip": "7ea06d5b0207ed3f20b5962bde7af18e6744fe089a98c51d66ac812252746342",
+    "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/pebble/com_github_cockroachdb_pebble-v0.0.0-20230807145728-40d3f411e45b.zip": "f0319e618ed024b7d708e2c1f8cf0d4b9b7e4112943288ba6036e893a7f8c151",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/redact/com_github_cockroachdb_redact-v1.1.5.zip": "11b30528eb0dafc8bc1a5ba39d81277c257cbe6946a7564402f588357c164560",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/returncheck/com_github_cockroachdb_returncheck-v0.0.0-20200612231554-92cdbca611dd.zip": "ce92ba4352deec995b1f2eecf16eba7f5d51f5aa245a1c362dfe24c83d31f82b",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/sentry-go/com_github_cockroachdb_sentry_go-v0.6.1-cockroachdb.2.zip": "fbb2207d02aecfdd411b1357efe1192dbb827959e36b7cab7491731ac55935c9",

go.mod

@@ -116,7 +116,7 @@ require (
 	github.com/cockroachdb/go-test-teamcity v0.0.0-20191211140407-cff980ad0a55
 	github.com/cockroachdb/gostdlib v1.19.0
 	github.com/cockroachdb/logtags v0.0.0-20230118201751-21c54148d20b
-	github.com/cockroachdb/pebble v0.0.0-20230728153158-ce43e6535942
+	github.com/cockroachdb/pebble v0.0.0-20230807145728-40d3f411e45b
 	github.com/cockroachdb/redact v1.1.5
 	github.com/cockroachdb/returncheck v0.0.0-20200612231554-92cdbca611dd
 	github.com/cockroachdb/stress v0.0.0-20220803192808-1806698b1b7b

go.sum

@@ -493,8 +493,8 @@ github.com/cockroachdb/gostdlib v1.19.0/go.mod h1:+dqqpARXbE/gRDEhCak6dm0l14AaTy
 github.com/cockroachdb/logtags v0.0.0-20211118104740-dabe8e521a4f/go.mod h1:Vz9DsVWQQhf3vs21MhPMZpMGSht7O/2vFW2xusFUVOs=
 github.com/cockroachdb/logtags v0.0.0-20230118201751-21c54148d20b h1:r6VH0faHjZeQy818SGhaone5OnYfxFR/+AzdY3sf5aE=
 github.com/cockroachdb/logtags v0.0.0-20230118201751-21c54148d20b/go.mod h1:Vz9DsVWQQhf3vs21MhPMZpMGSht7O/2vFW2xusFUVOs=
-github.com/cockroachdb/pebble v0.0.0-20230728153158-ce43e6535942 h1:y+9IilMcEy6qyDIZ94Vw3hJTk17J4j21TaMIe8C29Zg=
-github.com/cockroachdb/pebble v0.0.0-20230728153158-ce43e6535942/go.mod h1:FN5O47SBEz5+kO9fG8UTR64g2WS1u5ZFCgTvxGjoSks=
+github.com/cockroachdb/pebble v0.0.0-20230807145728-40d3f411e45b h1:ymYyDZy5WYRTBqPVYgy0XUW+gHx2HeRkyt1FJmNJVOo=
+github.com/cockroachdb/pebble v0.0.0-20230807145728-40d3f411e45b/go.mod h1:FN5O47SBEz5+kO9fG8UTR64g2WS1u5ZFCgTvxGjoSks=
 github.com/cockroachdb/redact v1.1.3/go.mod h1:BVNblN9mBWFyMyqK1k3AAiSxhvhfK2oOZZ2lK+dpvRg=
 github.com/cockroachdb/redact v1.1.5 h1:u1PMllDkdFfPWaNGMyLD1+so+aq3uUItthCFqzwPJ30=
 github.com/cockroachdb/redact v1.1.5/go.mod h1:BVNblN9mBWFyMyqK1k3AAiSxhvhfK2oOZZ2lK+dpvRg=

pkg/ccl/sqlproxyccl/BUILD.bazel

@@ -53,6 +53,7 @@ go_library(
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//credentials/insecure",
         "@org_golang_google_grpc//status",
+        "@org_golang_x_exp//slices",
     ],
 )
 
@@ -93,6 +94,7 @@ go_test(
         "//pkg/sql",
         "//pkg/sql/catalog/descs",
         "//pkg/sql/pgwire",
+        "//pkg/sql/pgwire/pgcode",
         "//pkg/sql/pgwire/pgerror",
         "//pkg/sql/pgwire/pgwirebase",
         "//pkg/testutils",

pkg/ccl/sqlproxyccl/authentication.go

@@ -13,10 +13,18 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/interceptor"
 	"github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/throttler"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/errors"
 	pgproto3 "github.com/jackc/pgproto3/v2"
+	"golang.org/x/exp/slices"
 )
 
+// These errors occur after successful auth but are sent back as a result of
+// the auth request. Receiving such error shouldn't trigger throttle.
+var pgPostAuthErrorCodes = []string{
+	pgcode.TooManyConnections.String(),
+}
+
 // authenticate handles the startup of the pgwire protocol to the point where
 // the connections is considered authenticated. If that doesn't happen, it
 // returns an error.
@@ -111,13 +119,21 @@ var authenticate = func(
 		// Server has rejected the authentication response from the client and
 		// has closed the connection.
 		case *pgproto3.ErrorResponse:
-			throttleError := throttleHook(throttler.AttemptInvalidCredentials)
+			// The error may be in response of auth message but may not indicate
+			// unsuccessful authentication. Clear throttle if this is the case.
+			var throttleError error
+			if slices.Contains(pgPostAuthErrorCodes, tp.Code) {
+				throttleError = throttleHook(throttler.AttemptOK)
+			} else {
+				throttleError = throttleHook(throttler.AttemptInvalidCredentials)
+			}
 			if throttleError != nil {
 				if err = feSend(toPgError(throttleError)); err != nil {
 					return nil, err
 				}
 				return nil, throttleError
 			}
+
 			if err = feSend(backendMsg); err != nil {
 				return nil, err
 			}

pkg/ccl/sqlproxyccl/authentication_test.go

@@ -13,6 +13,7 @@ import (
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/throttler"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/jackc/pgproto3/v2"
 	"github.com/stretchr/testify/assert"
@@ -98,14 +99,14 @@ func TestAuthenticateClearText(t *testing.T) {
 func TestAuthenticateThrottled(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 
-	server := func(t *testing.T, be *pgproto3.Backend, authResponse pgproto3.BackendMessage) {
+	server := func(t *testing.T, be *pgproto3.Backend) {
 		require.NoError(t, be.Send(&pgproto3.AuthenticationCleartextPassword{}))
 
 		msg, err := be.Receive()
 		require.NoError(t, err)
 		require.Equal(t, msg, &pgproto3.PasswordMessage{Password: "password"})
 
-		require.NoError(t, be.Send(authResponse))
+		require.NoError(t, be.Send(&pgproto3.ErrorResponse{Message: "wrong password"}))
 	}
 
 	client := func(t *testing.T, fe *pgproto3.Frontend) {
@@ -130,44 +131,79 @@ func TestAuthenticateThrottled(t *testing.T) {
 		require.Error(t, err)
 	}
 
-	type testCase struct {
-		name           string
-		result         pgproto3.BackendMessage
-		expectedStatus throttler.AttemptStatus
-	}
-	for _, tc := range []testCase{
-		{
-			name:           "AuthenticationOkay",
-			result:         &pgproto3.AuthenticationOk{},
-			expectedStatus: throttler.AttemptOK,
-		},
-		{
-			name:           "AuthenticationError",
-			result:         &pgproto3.ErrorResponse{Message: "wrong password"},
-			expectedStatus: throttler.AttemptInvalidCredentials,
-		},
-	} {
-		t.Run(tc.name, func(t *testing.T) {
-			proxyToServer, serverToProxy := net.Pipe()
-			proxyToClient, clientToProxy := net.Pipe()
-			sqlServer := pgproto3.NewBackend(pgproto3.NewChunkReader(serverToProxy), serverToProxy)
-			sqlClient := pgproto3.NewFrontend(pgproto3.NewChunkReader(clientToProxy), clientToProxy)
-
-			go server(t, sqlServer, &pgproto3.AuthenticationOk{})
-			go client(t, sqlClient)
-
-			_, err := authenticate(proxyToClient, proxyToServer, nil, /* proxyBackendKeyData */
-				func(status throttler.AttemptStatus) error {
-					require.Equal(t, throttler.AttemptOK, status)
-					return throttledError
-				})
-			require.Error(t, err)
-			require.Contains(t, err.Error(), "connection attempt throttled")
-
-			proxyToServer.Close()
-			proxyToClient.Close()
+	proxyToServer, serverToProxy := net.Pipe()
+	proxyToClient, clientToProxy := net.Pipe()
+	sqlServer := pgproto3.NewBackend(pgproto3.NewChunkReader(serverToProxy), serverToProxy)
+	sqlClient := pgproto3.NewFrontend(pgproto3.NewChunkReader(clientToProxy), clientToProxy)
+
+	go server(t, sqlServer)
+	go client(t, sqlClient)
+
+	_, err := authenticate(proxyToClient, proxyToServer, nil, /* proxyBackendKeyData */
+		func(status throttler.AttemptStatus) error {
+			require.Equal(t, throttler.AttemptInvalidCredentials, status)
+			return throttledError
 		})
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "connection attempt throttled")
+
+	proxyToServer.Close()
+	proxyToClient.Close()
+}
+
+func TestErrorFollowingAuthenticateNotThrottled(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	server := func(t *testing.T, be *pgproto3.Backend) {
+		require.NoError(t, be.Send(&pgproto3.AuthenticationCleartextPassword{}))
+
+		msg, err := be.Receive()
+		require.NoError(t, err)
+		require.Equal(t, msg, &pgproto3.PasswordMessage{Password: "password"})
+
+		require.NoError(t, be.Send(&pgproto3.ErrorResponse{
+			Code:    pgcode.TooManyConnections.String(),
+			Message: "sorry, too many clients already"}))
 	}
+
+	client := func(t *testing.T, fe *pgproto3.Frontend) {
+		msg, err := fe.Receive()
+		require.NoError(t, err)
+		require.Equal(t, msg, &pgproto3.AuthenticationCleartextPassword{})
+
+		require.NoError(t, fe.Send(&pgproto3.PasswordMessage{Password: "password"}))
+
+		// Try reading from the connection. This check ensures authorize
+		// swallowed the OK/Error response from the sql server.
+		msg, err = fe.Receive()
+		require.NoError(t, err)
+		require.Equal(t, msg, &pgproto3.ErrorResponse{
+			Code:    pgcode.TooManyConnections.String(),
+			Message: "sorry, too many clients already"})
+	}
+
+	proxyToServer, serverToProxy := net.Pipe()
+	proxyToClient, clientToProxy := net.Pipe()
+	sqlServer := pgproto3.NewBackend(pgproto3.NewChunkReader(serverToProxy), serverToProxy)
+	sqlClient := pgproto3.NewFrontend(pgproto3.NewChunkReader(clientToProxy), clientToProxy)
+
+	go server(t, sqlServer)
+	go client(t, sqlClient)
+
+	var throttleCallCount int
+	_, err := authenticate(proxyToClient, proxyToServer, nil, /* proxyBackendKeyData */
+		func(status throttler.AttemptStatus) error {
+			throttleCallCount++
+			require.Equal(t, throttler.AttemptOK, status)
+			return nil
+		})
+	require.Equal(t, 1, throttleCallCount)
+	require.Error(t, err)
+	require.Contains(t, err.Error(),
+		"codeAuthFailed: authentication failed: sorry, too many clients already")
+
+	proxyToServer.Close()
+	proxyToClient.Close()
 }
 
 func TestAuthenticateError(t *testing.T) {

pkg/sql/pgwire/conn.go

@@ -207,6 +207,9 @@ func (c *conn) processCommands(
 
 	var decrementConnectionCount func()
 	if decrementConnectionCount, retErr = sqlServer.IncrementConnectionCount(c.sessionArgs); retErr != nil {
+		// This will return pgcode.TooManyConnections which is used by the sql proxy
+		// to skip failed auth throttle (as in this case the auth was fine but the
+		// error occurred before sending back auth ok msg)
 		_ = c.sendError(ctx, retErr)
 		return
 	}