diff --git a/conn.go b/conn.go index 26a545d9..e9b7e736 100644 --- a/conn.go +++ b/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 { diff --git a/conn_test.go b/conn_test.go index ba8e0008..24263c97 100644 --- a/conn_test.go +++ b/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()) + } +} diff --git a/rows.go b/rows.go index 10c1dbd7..74664f6d 100644 --- a/rows.go +++ b/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: