diff --git a/buf_test.go b/buf_test.go new file mode 100644 index 00000000..37e55ca0 --- /dev/null +++ b/buf_test.go @@ -0,0 +1,182 @@ +package pq + +import ( + "bufio" + "bytes" + "database/sql/driver" + "encoding/binary" + "errors" + "math" + "testing" + + "github.com/lib/pq/internal/proto" + "github.com/lib/pq/oid" +) + +func TestRecvMessageInvalidFrameLength(t *testing.T) { + tests := []struct { + name string + code proto.ResponseCode + length int32 + wantErrMsg string + }{ + { + name: "zero length", + code: proto.DataRow, + length: 0, + wantErrMsg: "pq: lost synchronization with server: got message type \"(D) DataRow\", length 0", + }, + { + name: "length 1", + code: proto.DataRow, + length: 1, + wantErrMsg: "pq: lost synchronization with server: got message type \"(D) DataRow\", length 1", + }, + { + name: "length 2", + code: proto.DataRow, + length: 2, + wantErrMsg: "pq: lost synchronization with server: got message type \"(D) DataRow\", length 2", + }, + { + name: "length 3", + code: proto.DataRow, + length: 3, + wantErrMsg: "pq: lost synchronization with server: got message type \"(D) DataRow\", length 3", + }, + { + name: "negative length -1", + code: proto.DataRow, + length: -1, + wantErrMsg: "pq: lost synchronization with server: got message type \"(D) DataRow\", length -1", + }, + { + name: "negative length -4", + code: proto.DataRow, + length: -4, + wantErrMsg: "pq: lost synchronization with server: got message type \"(D) DataRow\", length -4", + }, + { + name: "negative length -100", + code: proto.DataRow, + length: -100, + wantErrMsg: "pq: lost synchronization with server: got message type \"(D) DataRow\", length -100", + }, + { + name: "negative length min int32", + code: proto.DataRow, + length: math.MinInt32, + wantErrMsg: "pq: lost synchronization with server: got message type \"(D) DataRow\", length -2147483648", + }, + { + name: "exceeds 1gb bound", + code: proto.DataRow, + length: 0x40000000, + wantErrMsg: "pq: lost synchronization with server: got message type \"(D) DataRow\", length 1073741824", + }, + { + name: "command complete zero length", + code: proto.CommandComplete, + length: 0, + wantErrMsg: "pq: lost synchronization with server: got message type \"(C) CommandComplete\", length 0", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + raw := make([]byte, 5) + raw[0] = byte(tt.code) + binary.BigEndian.PutUint32(raw[1:], uint32(tt.length)) + + cn := &conn{ + buf: bufio.NewReader(bytes.NewReader(raw)), + } + + var rb readBuf + var panicVal any + func() { + defer func() { + panicVal = recover() + }() + _, err := cn.recvMessage(&rb) + if err == nil { + t.Fatalf("expected error for length %d, got nil", tt.length) + } + if err.Error() != tt.wantErrMsg { + t.Errorf("wrong error message:\nhave: %s\nwant: %s", err.Error(), tt.wantErrMsg) + } + }() + + if panicVal != nil { + t.Fatalf("unexpected panic: %v", panicVal) + } + if err := cn.err.get(); !errors.Is(err, driver.ErrBadConn) { + t.Errorf("connection error not set to ErrBadConn, got %v", err) + } + }) + } +} + +func TestReadBuf(t *testing.T) { + t.Run("basic types", func(t *testing.T) { + data := []byte{ + 0x00, 0x00, 0x00, 0x2a, // int32: 42 + 0x00, 0x00, 0x04, 0xd2, // oid: 1234 + 0x00, 0x10, // int16: 16 + 'h', 'e', 'l', 'l', 'o', 0x00, // string: "hello" + 0x01, 0x02, 0x03, // next(3) + 0x7f, // byte + } + + b := readBuf(data) + if v := b.int32(); v != 42 { + t.Errorf("int32: have %d, want 42", v) + } + if v := b.oid(); v != oid.Oid(1234) { + t.Errorf("oid: have %d, want 1234", v) + } + if v := b.int16(); v != 16 { + t.Errorf("int16: have %d, want 16", v) + } + if v := b.string(); v != "hello" { + t.Errorf("string: have %q, want hello", v) + } + if v := b.next(3); !bytes.Equal(v, []byte{0x01, 0x02, 0x03}) { + t.Errorf("next: have %v, want [1 2 3]", v) + } + if v := b.byte(); v != 0x7f { + t.Errorf("byte: have 0x%x, want 0x7f", v) + } + if len(b) != 0 { + t.Errorf("expected empty buffer, remaining: %d", len(b)) + } + }) + + t.Run("string without null terminator panics", func(t *testing.T) { + b := readBuf([]byte("unterminated")) + defer func() { + r := recover() + if r == nil { + t.Fatal("expected panic for unterminated string") + } + }() + b.string() + }) +} + +func TestWriteBuf(t *testing.T) { + var wb writeBuf + wb.byte(proto.Query) + wb.int32(100) + wb.int16(20) + wb.string("test") + wb.bytes([]byte{1, 2}) + + if wb.buf[0] != byte(proto.Query) { + t.Errorf("first byte: have %c, want %c", wb.buf[0], proto.Query) + } + res := wb.wrap() + if len(res) == 0 { + t.Fatal("wrap returned empty buffer") + } +} diff --git a/conn.go b/conn.go index 01b046f0..eb712720 100644 --- a/conn.go +++ b/conn.go @@ -1106,7 +1106,8 @@ func (cn *conn) recvMessage(r *readBuf) (proto.ResponseCode, error) { // Read the type and length of the message that follows. t := proto.ResponseCode(x[0]) - n := int(binary.BigEndian.Uint32(x[1:])) - 4 + msgLen := int(int32(binary.BigEndian.Uint32(x[1:]))) + n := msgLen - 4 // When PostgreSQL cannot start a backend (e.g., an external process limit), // it sends plain text like "Ecould not fork new process [..]", which @@ -1120,6 +1121,10 @@ func (cn *conn) recvMessage(r *readBuf) (proto.ResponseCode, error) { msg, _ := cn.buf.ReadString('\x00') return 0, fmt.Errorf("pq: server error: %s%s", string(x[1:]), strings.TrimSuffix(msg, "\x00")) } + if n < 0 || msgLen > 0x3fffffff { + cn.err.set(driver.ErrBadConn) + return 0, fmt.Errorf("pq: lost synchronization with server: got message type %q, length %d", t, msgLen) + } if !proto.ValidLongMessageType(t) && n > proto.MaxMsgLen { return 0, fmt.Errorf("pq: lost synchronization with server: got message type %q, length %d", t, n) } diff --git a/conn_test.go b/conn_test.go index ba8e0008..103351e6 100644 --- a/conn_test.go +++ b/conn_test.go @@ -1861,6 +1861,66 @@ func TestLargeMessage(t *testing.T) { } } +func TestMalformedFrameLength(t *testing.T) { + tests := []struct { + name string + raw []byte + wantErr string + }{ + { + name: "zero length frame", + raw: []byte{byte(proto.DataRow), 0, 0, 0, 0}, + wantErr: `pq: lost synchronization with server: got message type "(D) DataRow", length 0`, + }, + { + name: "length smaller than header", + raw: []byte{byte(proto.DataRow), 0, 0, 0, 2}, + wantErr: `pq: lost synchronization with server: got message type "(D) DataRow", length 2`, + }, + { + name: "negative length frame", + raw: []byte{byte(proto.DataRow), 0xff, 0xff, 0xff, 0xff}, + wantErr: `pq: lost synchronization with server: got message type "(D) DataRow", length -1`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + f := pqtest.NewFake(t, func(f pqtest.Fake, cn net.Conn) { + f.Startup(cn, nil) + for { + code, q, ok := f.ReadMsg(cn) + if !ok { + return + } + switch code { + case proto.Query: + switch q := string(q[:bytes.IndexByte(q, 0)]); { + case q == ";": + f.WriteMsg(cn, proto.EmptyQueryResponse, "") + f.WriteMsg(cn, proto.ReadyForQuery, "I") + default: + _, _ = cn.Write(tt.raw) + } + case proto.Terminate: + cn.Close() + return + } + } + }) + defer f.Close() + + db := pqtest.MustDB(t, f.DSN()) + var have string + err := db.QueryRow(`select 1`).Scan(&have) + if !pqtest.ErrorContains(err, tt.wantErr) { + t.Fatalf("wrong error:\nhave: %s\nwant: %s", err, tt.wantErr) + } + }) + } +} + // reading from circularConn yields content[:prefixLen] once, followed by // content[prefixLen:] over and over again. It never returns EOF. type circularConn struct { diff --git a/rows.go b/rows.go index 10c1dbd7..a5cfcd27 100644 --- a/rows.go +++ b/rows.go @@ -124,6 +124,10 @@ func (rs *rows) Next(dest []driver.Value) (resErr error) { dest[i] = nil continue } + if l < 0 || l > len(rs.rb) { + rs.cn.err.set(driver.ErrBadConn) + return fmt.Errorf("pq: invalid column length %d in DataRow", l) + } dest[i], err = decode(&rs.cn.parameterStatus, rs.rb.next(l), rs.colTyps[i].OID, rs.colFmts[i]) if err != nil { return rs.cn.handleError(err)