pgwire: delay initializing the SQL session until authentication succe…

…eeds.

This patch does what it says on the cover.

It also introduces a separate method `handleAuthentication()` on the
v3conn object to deal with authentication. This method should
overwrite the User in SessionArgs if the protocol exchange decides a
new username.

As an added benefit this patch also tags logging messages with the
authenticated user.

pkg/sql/pgwire/server.go

@@ -26,7 +26,6 @@ import (
 	"golang.org/x/net/context"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
-	"github.com/cockroachdb/cockroach/pkg/security"
 	"github.com/cockroachdb/cockroach/pkg/sql"
 	"github.com/cockroachdb/cockroach/pkg/util"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -222,7 +221,7 @@ func (s *Server) ServeConn(ctx context.Context, conn net.Conn) error {
 		// We make a connection regardless of argsErr. If there was an error parsing
 		// the args, the connection will only be used to send a report of that
 		// error.
-		v3conn := makeV3Conn(ctx, conn, s.executor, s.metrics, sessionArgs)
+		v3conn := makeV3Conn(ctx, conn, s.metrics)
 		defer v3conn.finish(ctx)
 		if argsErr != nil {
 			return v3conn.sendInternalError(argsErr.Error())
@@ -236,15 +235,10 @@ func (s *Server) ServeConn(ctx context.Context, conn net.Conn) error {
 			return v3conn.sendInternalError(ErrDraining)
 		}
 
-		if tlsConn, ok := conn.(*tls.Conn); ok {
-			tlsState := tlsConn.ConnectionState()
-			authenticationHook, err := security.UserAuthHook(s.cfg.Insecure, &tlsState)
-			if err != nil {
-				return v3conn.sendInternalError(err.Error())
-			}
-			return v3conn.serve(ctx, authenticationHook)
+		if err := v3conn.handleAuthentication(ctx, s.cfg.Insecure, &sessionArgs); err != nil {
+			return v3conn.sendInternalError(err.Error())
 		}
-		return v3conn.serve(ctx, nil)
+		return v3conn.serve(ctx, s.executor, sessionArgs)
 	}
 
 	return errors.Errorf("unknown protocol version %d", version)

pkg/sql/pgwire/v3.go

@@ -19,6 +19,7 @@ package pgwire
 
 import (
 	"bufio"
+	"crypto/tls"
 	"fmt"
 	"net"
 	"strconv"
@@ -27,6 +28,7 @@ import (
 	"github.com/pkg/errors"
 	"golang.org/x/net/context"
 
+	"github.com/cockroachdb/cockroach/pkg/security"
 	"github.com/cockroachdb/cockroach/pkg/sql"
 	"github.com/cockroachdb/cockroach/pkg/sql/parser"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
@@ -127,31 +129,25 @@ type v3Conn struct {
 	metrics ServerMetrics
 }
 
-func makeV3Conn(
-	ctx context.Context,
-	conn net.Conn,
-	executor *sql.Executor,
-	metrics ServerMetrics,
-	sessionArgs sql.SessionArgs,
-) v3Conn {
+func makeV3Conn(ctx context.Context, conn net.Conn, metrics ServerMetrics) v3Conn {
 	return v3Conn{
 		conn:     conn,
 		rd:       bufio.NewReader(conn),
 		wr:       bufio.NewWriter(conn),
-		executor: executor,
 		writeBuf: writeBuffer{bytecount: metrics.BytesOutCount},
 		metrics:  metrics,
-		session:  sql.NewSession(ctx, sessionArgs, executor, conn.RemoteAddr()),
 	}
 }
 
 func (c *v3Conn) finish(ctx context.Context) {
+	if c.session != nil {
+		c.session.Finish(c.executor)
+	}
 	// This is better than always flushing on error.
 	if err := c.wr.Flush(); err != nil {
 		log.Error(ctx, err)
 	}
 	_ = c.conn.Close()
-	c.session.Finish(c.executor)
 }
 
 func parseOptions(data []byte) (sql.SessionArgs, error) {
@@ -196,17 +192,37 @@ var statusReportParams = map[string]string{
 	"server_version": "9.5.0",
 }
 
-func (c *v3Conn) serve(ctx context.Context, authenticationHook func(string, bool) error) error {
-	if authenticationHook != nil {
-		if err := authenticationHook(c.session.User, true /* public */); err != nil {
+// handleAuthentication should discuss with the client to arrange
+// authentication and update sessionArgs with the authenticated user's
+// name, if different from the one given initially. Note: at this
+// point the sql.Session does not exist yet! If need exists to access the
+// database to look up authentication data, use the internal executor.
+func (c *v3Conn) handleAuthentication(
+	ctx context.Context, insecure bool, sessionArgs *sql.SessionArgs,
+) error {
+	if tlsConn, ok := c.conn.(*tls.Conn); ok {
+		tlsState := tlsConn.ConnectionState()
+		authenticationHook, err := security.UserAuthHook(insecure, &tlsState)
+		if err != nil {
 			return c.sendInternalError(err.Error())
 		}
+		if authenticationHook != nil {
+			if err := authenticationHook(sessionArgs.User, true /* public */); err != nil {
+				return c.sendInternalError(err.Error())
+			}
+		}
 	}
 	c.writeBuf.initMsg(serverMsgAuth)
 	c.writeBuf.putInt32(authOK)
 	if err := c.writeBuf.finishMsg(c.wr); err != nil {
 		return err
 	}
+	return nil
+}
+
+func (c *v3Conn) serve(
+	ctx context.Context, executor *sql.Executor, sessionArgs sql.SessionArgs,
+) error {
 	for key, value := range statusReportParams {
 		c.writeBuf.initMsg(serverMsgParameterStatus)
 		c.writeBuf.writeTerminatedString(key)
@@ -219,6 +235,10 @@ func (c *v3Conn) serve(ctx context.Context, authenticationHook func(string, bool
 		return err
 	}
 
+	ctx = log.WithLogTagStr(ctx, "user", sessionArgs.User)
+	c.executor = executor
+	c.session = sql.NewSession(ctx, sessionArgs, executor, c.conn.RemoteAddr())
+
 	for {
 		if !c.doingExtendedQueryMessage {
 			c.writeBuf.initMsg(serverMsgReady)

pkg/sql/pgwire/v3_test.go

@@ -32,9 +32,7 @@ func makeTestV3Conn(c net.Conn) v3Conn {
 	return makeV3Conn(
 		context.Background(),
 		c,
-		sql.NewDummyExecutor(),
 		makeServerMetrics(),
-		sql.SessionArgs{},
 	)
 }
 
@@ -90,5 +88,9 @@ func testMaliciousInput(t *testing.T, data []byte) {
 	}()
 
 	v3Conn := makeTestV3Conn(r)
-	_ = v3Conn.serve(context.Background(), nil)
+	_ = v3Conn.serve(
+		context.Background(),
+		sql.NewDummyExecutor(),
+		sql.SessionArgs{},
+	)
 }