Don't require CommandComplete or EmptyQueryResponse in simpleQuery()

CHANGELOG.md

@@ -59,6 +59,9 @@ newer. Previously PostgreSQL 8.4 and newer were supported.
 
 - Decode bpchar into a string ([#949]).
 
+- Fix panic in Ping() by not requiring CommandComplete or EmptyQueryResponse in
+  simpleQuery() ([#1234])
+
 - Recognize bit/varbit ([#743]) and float types ([#1166]) in ColumnTypeScanType().
 
 - Accept `PGGSSLIB` and `PGKRBSRVNAME` environment variables ([#1143]).
@@ -91,6 +94,7 @@ newer. Previously PostgreSQL 8.4 and newer were supported.
 [#1224]: https://github.com/lib/pq/pull/1224
 [#1226]: https://github.com/lib/pq/pull/1226
 [#1228]: https://github.com/lib/pq/pull/1228
+[#1234]: https://github.com/lib/pq/pull/1234
 
 
 v1.10.9 (2023-04-26)

conn.go

@@ -512,7 +512,7 @@ func (cn *conn) simpleExec(q string) (res driver.Result, commandTag string, resE
 	}
 }
 
-func (cn *conn) simpleQuery(q string) (res *rows, resErr error) {
+func (cn *conn) simpleQuery(q string) (*rows, error) {
 	if debugProto {
 		fmt.Fprintf(os.Stderr, "         START conn.simpleQuery\n")
 		defer fmt.Fprintf(os.Stderr, "         END conn.simpleQuery\n")
@@ -525,6 +525,10 @@ func (cn *conn) simpleQuery(q string) (res *rows, resErr error) {
 		return nil, cn.handleError(err, q)
 	}
 
+	var (
+		res    *rows
+		resErr error
+	)
 	for {
 		t, r, err := cn.recv1()
 		if err != nil {
@@ -533,9 +537,9 @@ func (cn *conn) simpleQuery(q string) (res *rows, resErr error) {
 		switch t {
 		case proto.CommandComplete, proto.EmptyQueryResponse:
 			// We allow queries which don't return any results through Query as
-			// well as Exec.  We still have to give database/sql a rows object
+			// well as Exec. We still have to give database/sql a rows object
 			// the user can close, though, to avoid connections from being
-			// leaked.  A "rows" with done=true works fine for that purpose.
+			// leaked. A "rows" with done=true works fine for that purpose.
 			if resErr != nil {
 				cn.err.set(driver.ErrBadConn)
 				return nil, fmt.Errorf("pq: unexpected message %q in simple query execution", t)
@@ -558,8 +562,10 @@ func (cn *conn) simpleQuery(q string) (res *rows, resErr error) {
 			res.done = true
 		case proto.ReadyForQuery:
 			cn.processReadyForQuery(r)
-			// done
-			return res, cn.handleError(resErr, q)
+			if err == nil && res == nil {
+				res = &rows{done: true}
+			}
+			return res, cn.handleError(resErr, q) // done
 		case proto.ErrorResponse:
 			res = nil
 			resErr = parseError(r, q)
@@ -568,13 +574,11 @@ func (cn *conn) simpleQuery(q string) (res *rows, resErr error) {
 				cn.err.set(driver.ErrBadConn)
 				return nil, fmt.Errorf("pq: unexpected DataRow in simple query execution")
 			}
-			// the query didn't fail; kick off to Next
-			return res, cn.saveMessage(t, r)
+			return res, cn.saveMessage(t, r) // The query didn't fail; kick off to Next
 		case proto.RowDescription:
 			// res might be non-nil here if we received a previous
-			// CommandComplete, but that's fine; just overwrite it
-			res = &rows{cn: cn}
-			res.rowsHeader = parsePortalRowDescribe(r)
+			// CommandComplete, but that's fine and just overwrite it.
+			res = &rows{cn: cn, rowsHeader: parsePortalRowDescribe(r)}
 
 			// To work around a bug in QueryRow in Go 1.2 and earlier, wait
 			// until the first DataRow has been received.

conn_test.go

@@ -17,6 +17,7 @@ import (
 
 	"github.com/lib/pq/internal/pgpass"
 	"github.com/lib/pq/internal/pqtest"
+	"github.com/lib/pq/internal/proto"
 )
 
 func TestReconnect(t *testing.T) {
@@ -904,6 +905,41 @@ func TestSimpleQuery(t *testing.T) {
 	}
 }
 
+// Make sure SimpleQuery doesn't panic if there is no query response. See #1059
+// and #1173
+func TestSimpleQueryWithoutResponse(t *testing.T) {
+	t.Parallel()
+
+	f := pqtest.NewFake(t)
+	f.Accept(func(cn net.Conn) {
+		f.Startup(cn)
+
+		for {
+			code, _, ok := f.ReadMsg(cn)
+			if !ok {
+				return
+			}
+			switch code {
+			case proto.Query:
+				// Make sure we DON'T send this
+				//f.WriteMsg(cn, proto.EmptyQueryResponse)
+				f.WriteMsg(cn, proto.ReadyForQuery, 'I')
+			case proto.Terminate:
+				cn.Close()
+				return
+			}
+		}
+	})
+
+	db := pqtest.MustDB(t, f.DSN())
+	if err := db.Ping(); err != nil {
+		t.Fatal(err)
+	}
+	if err := db.Ping(); err != nil {
+		t.Fatal(err)
+	}
+}
+
 func TestBindError(t *testing.T) {
 	t.Parallel()
 	db := pqtest.MustDB(t)

internal/pqtest/fake.go

@@ -0,0 +1,119 @@
+package pqtest
+
+import (
+	"encoding/binary"
+	"net"
+	"testing"
+
+	"github.com/lib/pq/internal/proto"
+)
+
+type Fake struct {
+	l net.Listener
+	t testing.TB
+}
+
+// NewFake creates a new "fake" PostgreSQL server. You need to accept
+// connections with [Fake.Accept].
+//
+// This can also be tested against libpq with something like:
+//
+//	f := pqtest.NewFake(t)
+//	f.Accept(..)
+//
+//	fmt.Println("\n" + f.DSN())
+//	time.Sleep(9 * time.Minute)
+func NewFake(t testing.TB) Fake {
+	l, err := net.Listen("tcp", "127.0.0.1:")
+	if err != nil {
+		t.Fatal(err)
+	}
+	return Fake{l: l, t: t}
+}
+
+// DSN is the DSN to connect to for this server.
+func (f Fake) DSN() string {
+	h, p, err := net.SplitHostPort(f.l.Addr().String())
+	if err != nil {
+		f.t.Fatal(err)
+	}
+	return "host=" + h + " port=" + p
+}
+
+// Accept callback for new connections.
+func (f Fake) Accept(fun func(net.Conn)) {
+	go func() {
+		for {
+			cn, err := f.l.Accept()
+			if err != nil {
+				f.t.Errorf("accepting connection: %s", err)
+				return
+			}
+			go fun(cn)
+		}
+	}()
+}
+
+// Startup reads the startup message from the server with [f.ReadStartup] and
+// sends [proto.AuthenticationRequest] and [proto.ReadyForQuery].
+func (f Fake) Startup(cn net.Conn) {
+	if !f.ReadStartup(cn) {
+		return
+	}
+	// Technically we don't *need* to send the AuthRequest, but the psql CLI
+	// expects it.
+	f.WriteMsg(cn, proto.AuthenticationRequest, 0, 0, 0, 0)
+	f.WriteMsg(cn, proto.ReadyForQuery, 'I')
+}
+
+// ReadStart reads the startup message.
+func (f Fake) ReadStartup(cn net.Conn) bool {
+	_, _, ok := f.read(cn, true)
+	return ok
+}
+
+// ReadMsg reads a message from the client (frontend).
+func (f Fake) ReadMsg(cn net.Conn) (proto.RequestCode, []byte, bool) {
+	return f.read(cn, false)
+}
+
+func (f Fake) read(cn net.Conn, startup bool) (proto.RequestCode, []byte, bool) {
+	// Startup message has no code and only a length (herp derp).
+	sz := 5
+	if startup {
+		sz = 4
+	}
+	typ := make([]byte, sz)
+	_, err := cn.Read(typ)
+	if err != nil {
+		f.t.Errorf("reading: %s", err)
+		return 0, nil, false
+	}
+
+	var (
+		code   = proto.RequestCode(typ[0])
+		length = typ[1:]
+	)
+	if startup {
+		code = 0
+		length = typ
+	}
+
+	data := make([]byte, int(binary.BigEndian.Uint32(length))-4)
+	_, err = cn.Read(data)
+	if err != nil {
+		f.t.Errorf("reading: %s", err)
+		return 0, nil, false
+	}
+	return code, data, true
+}
+
+// WriteMsg writes a message to the client (frontend).
+func (f Fake) WriteMsg(cn net.Conn, code proto.ResponseCode, msg ...byte) {
+	l := []byte{byte(code), 0, 0, 0, 0}
+	binary.BigEndian.PutUint32(l[1:], uint32(len(msg)+4))
+	_, err := cn.Write(append(l, msg...))
+	if err != nil {
+		f.t.Error(err)
+	}
+}