Merge pull request cockroachdb#36504 from andreimatei/backport2.0-35776

release-2.0: pgwire: cancel authentication work when the network connection is closed
ajkr · Apr 5, 2019 · 3d0cb04 · 3d0cb04
2 parents 3df579a + c52b761
commit 3d0cb04
Show file tree

Hide file tree

Showing 8 changed files with 465 additions and 126 deletions.
diff --git a/pkg/base/testing_knobs.go b/pkg/base/testing_knobs.go
@@ -28,6 +28,7 @@ type TestingKnobs struct {
 	SQLExecutor         ModuleTestingKnobs
 	SQLLeaseManager     ModuleTestingKnobs
 	SQLSchemaChanger    ModuleTestingKnobs
+	PGWireTestingKnobs  ModuleTestingKnobs
 	SQLMigrationManager ModuleTestingKnobs
 	DistSQL             ModuleTestingKnobs
 	SQLEvalContext      ModuleTestingKnobs

diff --git a/pkg/server/server.go b/pkg/server/server.go
@@ -33,15 +33,6 @@ import (
 	"sync/atomic"
 	"time"
 
-	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
-
-	"github.com/elazarl/go-bindata-assetfs"
-	raven "github.com/getsentry/raven-go"
-	gwruntime "github.com/grpc-ecosystem/grpc-gateway/runtime"
-	opentracing "github.com/opentracing/opentracing-go"
-	"github.com/pkg/errors"
-	"google.golang.org/grpc"
-
 	"github.com/cockroachdb/cmux"
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/gossip"
@@ -59,6 +50,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/distsqlrun"
 	"github.com/cockroachdb/cockroach/pkg/sql/jobs"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sqlmigrations"
 	"github.com/cockroachdb/cockroach/pkg/storage"
 	"github.com/cockroachdb/cockroach/pkg/storage/engine"
@@ -78,6 +70,12 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
+	assetfs "github.com/elazarl/go-bindata-assetfs"
+	raven "github.com/getsentry/raven-go"
+	gwruntime "github.com/grpc-ecosystem/grpc-gateway/runtime"
+	opentracing "github.com/opentracing/opentracing-go"
+	"github.com/pkg/errors"
+	"google.golang.org/grpc"
 )
 
 var (
@@ -555,6 +553,9 @@ func NewServer(cfg Config, stopper *stop.Stopper) (*Server, error) {
 	if s.cfg.UseLegacyConnHandling {
 		s.registry.AddMetricStruct(s.sqlExecutor)
 	}
+	if pgwireKnobs := s.cfg.TestingKnobs.PGWireTestingKnobs; pgwireKnobs != nil {
+		execCfg.PGWireTestingKnobs = pgwireKnobs.(*sql.PGWireTestingKnobs)
+	}
 
 	s.pgServer = pgwire.MakeServer(
 		s.cfg.AmbientCtx,
@@ -638,8 +639,7 @@ type ListenError struct {
 func inspectEngines(
 	ctx context.Context,
 	engines []engine.Engine,
-	minVersion,
-	serverVersion roachpb.Version,
+	minVersion, serverVersion roachpb.Version,
 	clusterIDContainer *base.ClusterIDContainer,
 ) (
 	bootstrappedEngines []engine.Engine,
@@ -1506,7 +1506,7 @@ If problems persist, please see ` + base.DocsURL("cluster-setup-troubleshooting.
 			} else {
 				serveFn = s.pgServer.ServeConn
 			}
-			if err := serveFn(connCtx, conn); err != nil && !netutil.IsClosedConnection(err) {
+			if err := serveFn(connCtx, conn); err != nil {
 				// Report the error on this connection's context, so that we
 				// know which remote client caused the error when looking at
 				// the logs.
@@ -1538,11 +1538,7 @@ If problems persist, please see ` + base.DocsURL("cluster-setup-troubleshooting.
 			}
 			netutil.FatalIfUnexpected(httpServer.ServeWith(pgCtx, s.stopper, unixLn, func(conn net.Conn) {
 				connCtx := log.WithLogTagStr(pgCtx, "client", conn.RemoteAddr().String())
-				if err := s.pgServer.ServeConn(connCtx, conn); err != nil &&
-					!netutil.IsClosedConnection(err) {
-					// Report the error on this connection's context, so that we
-					// know which remote client caused the error when looking at
-					// the logs.
+				if err := s.pgServer.ServeConn(connCtx, conn); err != nil {
 					log.Error(connCtx, err)
 				}
 			}))
@@ -1640,7 +1636,7 @@ func (s *Server) doDrain(
 // On failure, the system may be in a partially drained state and should be
 // recovered by calling Undrain() with the same (or a larger) slice of modes.
 func (s *Server) Drain(ctx context.Context, on []serverpb.DrainMode) ([]serverpb.DrainMode, error) {
-	return s.doDrain(ctx, on, true)
+	return s.doDrain(ctx, on, true /* setTo */)
 }
 
 // Undrain idempotently deactivates the given DrainModes on the Server in the

diff --git a/pkg/sql/conn_executor.go b/pkg/sql/conn_executor.go
@@ -23,9 +23,6 @@ import (
 	"time"
 	"unicode/utf8"
 
-	"github.com/pkg/errors"
-	"golang.org/x/net/trace"
-
 	"github.com/cockroachdb/cockroach/pkg/config"
 	"github.com/cockroachdb/cockroach/pkg/internal/client"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -46,6 +43,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/uint128"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
+	"github.com/pkg/errors"
+	"golang.org/x/net/trace"
 )
 
 // A connExecutor is in charge of executing queries received on a given client
@@ -907,8 +906,8 @@ func (ex *connExecutor) run(ctx context.Context) error {
 			ex.phaseTimes[sessionStartParse] = tcmd.ParseStart
 			ex.phaseTimes[sessionEndParse] = tcmd.ParseEnd
 
-			ctx := withStatement(ex.Ctx(), ex.curStmt)
-			ev, payload, err = ex.execStmt(ctx, curStmt, stmtRes, nil /* pinfo */, pos)
+			stmtCtx := withStatement(ex.Ctx(), ex.curStmt)
+			ev, payload, err = ex.execStmt(stmtCtx, curStmt, stmtRes, nil /* pinfo */, pos)
 			if err != nil {
 				return err
 			}
@@ -960,16 +959,16 @@ func (ex *connExecutor) run(ctx context.Context) error {
 				ExpectedTypes: portal.Stmt.Columns,
 				AnonymizedStr: portal.Stmt.AnonymizedStr,
 			}
-			ctx := withStatement(ex.Ctx(), ex.curStmt)
-			ev, payload, err = ex.execStmt(ctx, curStmt, stmtRes, pinfo, pos)
+			stmtCtx := withStatement(ex.Ctx(), ex.curStmt)
+			ev, payload, err = ex.execStmt(stmtCtx, curStmt, stmtRes, pinfo, pos)
 			if err != nil {
 				return err
 			}
 		case PrepareStmt:
 			ex.curStmt = tcmd.Stmt
 			res = ex.clientComm.CreatePrepareResult(pos)
-			ctx := withStatement(ex.Ctx(), ex.curStmt)
-			ev, payload = ex.execPrepare(ctx, tcmd)
+			stmtCtx := withStatement(ex.Ctx(), ex.curStmt)
+			ev, payload = ex.execPrepare(stmtCtx, tcmd)
 		case DescribeStmt:
 			descRes := ex.clientComm.CreateDescribeResult(pos)
 			res = descRes

diff --git a/pkg/sql/executor.go b/pkg/sql/executor.go
@@ -24,8 +24,6 @@ import (
 	"sync"
 	"time"
 
-	"github.com/pkg/errors"
-
 	"github.com/cockroachdb/apd"
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/config"
@@ -56,6 +54,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/uint128"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
+	"github.com/pkg/errors"
 )
 
 // ClusterOrganization is the organization name.
@@ -271,6 +270,7 @@ type ExecutorConfig struct {
 	AuditLogger     *log.SecondaryLogger
 
 	TestingKnobs              *ExecutorTestingKnobs
+	PGWireTestingKnobs        *PGWireTestingKnobs
 	SchemaChangerTestingKnobs *SchemaChangerTestingKnobs
 	EvalContextTestingKnobs   tree.EvalContextTestingKnobs
 	// HistogramWindowInterval is (server.Config).HistogramWindowInterval.
@@ -296,6 +296,18 @@ var _ base.ModuleTestingKnobs = &ExecutorTestingKnobs{}
 // ModuleTestingKnobs is part of the base.ModuleTestingKnobs interface.
 func (*ExecutorTestingKnobs) ModuleTestingKnobs() {}
 
+// PGWireTestingKnobs contains knobs for the pgwire module.
+type PGWireTestingKnobs struct {
+	// AuthHook is used to override the normal authentication handling on new
+	// connections.
+	AuthHook func(context.Context) error
+}
+
+var _ base.ModuleTestingKnobs = &PGWireTestingKnobs{}
+
+// ModuleTestingKnobs implements the base.ModuleTestingKnobs interface.
+func (*PGWireTestingKnobs) ModuleTestingKnobs() {}
+
 // StatementFilter is the type of callback that
 // ExecutorTestingKnobs.StatementFilter takes.
 type StatementFilter func(context.Context, string, error)