Return an error instead of panicking on unknown response format code

conn.go

@@ -717,7 +717,11 @@ func (cn *conn) simpleQuery(q string) (*rows, error) {
 		case proto.RowDescription:
 			// res might be non-nil here if we received a previous
 			// CommandComplete, but that's fine and just overwrite it.
-			res = &rows{cn: cn, rowsHeader: parsePortalRowDescribe(r)}
+			rh, err := parsePortalRowDescribe(r)
+			if err != nil {
+				return nil, cn.handleError(err, q)
+			}
+			res = &rows{cn: cn, rowsHeader: rh}
 
 			// To work around a bug in QueryRow in Go 1.2 and earlier, wait
 			// until the first DataRow has been received.
@@ -1792,7 +1796,7 @@ func (cn *conn) readPortalDescribeResponse() (rowsHeader, error) {
 	}
 	switch t {
 	case proto.RowDescription:
-		return parsePortalRowDescribe(r), nil
+		return parsePortalRowDescribe(r)
 	case proto.NoData:
 		return rowsHeader{}, nil
 	case proto.ErrorResponse:
@@ -1910,7 +1914,7 @@ func parseStatementRowDescribe(r *readBuf) (colNames []string, colTyps []fieldDe
 	return
 }
 
-func parsePortalRowDescribe(r *readBuf) rowsHeader {
+func parsePortalRowDescribe(r *readBuf) (rowsHeader, error) {
 	n := r.int16()
 	colNames := make([]string, n)
 	colFmts := make([]format, n)
@@ -1921,13 +1925,17 @@ func parsePortalRowDescribe(r *readBuf) rowsHeader {
 		colTyps[i].OID = r.oid()
 		colTyps[i].Len = r.int16()
 		colTyps[i].Mod = r.int32()
-		colFmts[i] = format(r.int16())
+		f := format(r.int16())
+		if f != formatText && f != formatBinary {
+			return rowsHeader{}, fmt.Errorf("pq: unknown response format code: %d", f)
+		}
+		colFmts[i] = f
 	}
 	return rowsHeader{
 		colNames: colNames,
 		colFmts:  colFmts,
 		colTyps:  colTyps,
-	}
+	}, nil
 }
 
 func (cn *conn) ResetSession(ctx context.Context) error {

conn_test.go

@@ -2100,3 +2100,35 @@ func BenchmarkPreparedSelect(b *testing.B) {
 		})
 	})
 }
+
+func TestParsePortalRowDescribeInvalidFormat(t *testing.T) {
+	// Construct a fake RowDescription wire payload
+	// 1 column (n = 1)
+	// name "col1" (5 bytes: 'c','o','l','1', 0)
+	// next(6) skips 6 bytes
+	// OID (4 bytes)
+	// Len (2 bytes)
+	// Mod (4 bytes)
+	// format code = 2 (invalid) (2 bytes)
+	buf := []byte{
+		0, 1, // n = 1
+		'c', 'o', 'l', '1', 0, // name
+		0, 0, 0, 0, 0, 0, // next(6)
+		0, 0, 0, 0, // oid
+		0, 0,       // len
+		0, 0, 0, 0, // mod
+		0, 2,       // format = 2 (invalid)
+	}
+	r := readBuf(buf)
+
+	_, err := parsePortalRowDescribe(&r)
+
+	if err == nil {
+		t.Fatalf("expected error for invalid format code, got nil")
+	}
+
+	expectedMsg := "pq: unknown response format code: 2"
+	if err.Error() != expectedMsg {
+		t.Fatalf("expected error %q, got %q", expectedMsg, err.Error())
+	}
+}

rows.go

@@ -131,7 +131,10 @@ func (rs *rows) Next(dest []driver.Value) (resErr error) {
 			}
 			return rs.cn.handleError(resErr)
 		case proto.RowDescription:
-			next := parsePortalRowDescribe(&rs.rb)
+			next, err := parsePortalRowDescribe(&rs.rb)
+			if err != nil {
+				return rs.cn.handleError(err)
+			}
 			rs.next = &next
 			return io.EOF
 		default: