sqlproxyccl: fix non-auth error causing throttle

Currently an error that occurs after a successful authentication but
before the SQL server is ready to serve queries is considered as
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 sucessful as
conditions to throttle.

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

Release note: None

Author: darinpp

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/sql/tests",

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

@@ -210,6 +210,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
 	}