diff --git a/ctl/server.go b/ctl/server.go index 6dc4eb0a6..378f951a3 100644 --- a/ctl/server.go +++ b/ctl/server.go @@ -88,4 +88,7 @@ func BuildServerFlags(cmd *cobra.Command, srv *server.Command) { // Transactional storage engine flags.StringVarP(&srv.Config.Txsrc, "tx", "", "", "transaction/storage to use: one of roaring, rbf, badger, rbf_roaring, roaring_rbf, badger_roaring, roaring_badger, badger_rbf, or rbf_badger (default roaring)") + + // Postgres endpoint + flags.StringVar(&srv.Config.Postgres.Addr, "postgres.addr", "", "address to which to bind a postgres endpoint") } diff --git a/go.mod b/go.mod index 70221964e..577efd92f 100644 --- a/go.mod +++ b/go.mod @@ -20,6 +20,7 @@ require ( github.com/gorilla/handlers v1.3.0 github.com/gorilla/mux v1.7.0 github.com/hashicorp/memberlist v0.1.3 + github.com/lib/pq v1.8.0 github.com/opentracing/opentracing-go v1.1.0 github.com/pelletier/go-toml v1.2.0 github.com/pkg/errors v0.8.1 diff --git a/go.sum b/go.sum index f97f5d051..10c1da24b 100644 --- a/go.sum +++ b/go.sum @@ -117,6 +117,8 @@ github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORN github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= +github.com/lib/pq v1.8.0 h1:9xohqzkUwzR4Ga4ivdTcawVS89YSDVxXMa3xJX3cGzg= +github.com/lib/pq v1.8.0/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o= github.com/magiconair/properties v1.8.0 h1:LLgXmsheXeRoUOBOjtwPQCWIYqM/LU1ayDtDePerRcY= github.com/magiconair/properties v1.8.0/go.mod h1:PppfXfuXeibc/6YijjN8zIbojt8czPbwD3XqdrwzmxQ= github.com/matttproud/golang_protobuf_extensions v1.0.1 h1:4hp9jkHxhMHkqkrB3Ix0jegS5sx/RkqARlsWZ6pIwiU= diff --git a/pg/io.go b/pg/io.go new file mode 100644 index 000000000..80ae8cd24 --- /dev/null +++ b/pg/io.go @@ -0,0 +1,249 @@ +// Copyright 2020 Pilosa Corp. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pg + +import ( + "net" + "sync" + "sync/atomic" + "time" + "unsafe" + + "github.com/pkg/errors" +) + +// timeoutWriter wraps a connection and implements io.Writer with a timeout for each write. +type timeoutWriter struct { + conn net.Conn + timeout time.Duration +} + +func (w *timeoutWriter) Write(data []byte) (int, error) { + err := w.conn.SetWriteDeadline(time.Now().Add(w.timeout)) + if err != nil { + return 0, err + } + return w.conn.Write(data) +} + +// errPreempted is an error used to indicate preemption of an idle connection. +var errPreempted = errors.New("preempted during idle") + +// idleState is an atomic state value used to track a preemptible connection. +type idleState uint32 + +const ( + idleStateActive idleState = iota + idleStateIdle + idleStatePreempted + idleStatePendingPreemption +) + +func (s *idleState) load() idleState { + return idleState(atomic.LoadUint32((*uint32)(unsafe.Pointer(s)))) +} + +func (s *idleState) cas(old, new idleState) bool { + return atomic.CompareAndSwapUint32((*uint32)(unsafe.Pointer(s)), uint32(old), uint32(new)) +} + +// idleReader is an io.Reader implementation on a preemptible network connection. +// The connection has 2 modes: "idle" and "active". +// While in idle mode, the connection has no timeout but can be preempted. +// While in active mode, the connection may have a read timeout but cannot be immediately preempted. +// When a read completes in idle mode, the connection returns to active mode. +// If the connection is preempted in active mode, the preemption will be deferred until the connection returns to idle mode. +// This also provides a read timeout. +type idleReader struct { + conn net.Conn + timeout time.Duration + state idleState + preemptMu sync.Mutex +} + +// setIdle pushes the reader into idle mode. +// If a preemption is pending, it will be delivered on the next call to Read. +func (r *idleReader) setIdle() error { + // Clear the read deadline. + err := r.conn.SetReadDeadline(time.Time{}) + if err != nil { + return errors.Wrap(err, "failed to clear deadline") + } + + for { + // Transition to idle mode. + state := r.state.load() + var target idleState + switch state { + case idleStateActive: + // active -> idle + target = idleStateIdle + case idleStatePendingPreemption: + // pending preemption -> preempted + // Switching to idle mode activates the preemption. + target = idleStatePreempted + default: + panic("inconsistent state") + } + + if r.state.cas(state, target) { + return nil + } + } +} + +// preempt the reader. +// If the reader is not currently idle, the preemption will be delivered next time the connection enters idle mode. +// This does not wait until the preemption error is delivered. +func (r *idleReader) preempt() error { + r.preemptMu.Lock() + defer r.preemptMu.Unlock() + for { + state := r.state.load() + var target idleState + switch state { + case idleStateActive: + // active -> pending preemption + target = idleStatePendingPreemption + case idleStateIdle: + // idle -> preempted + target = idleStatePreempted + case idleStatePendingPreemption, idleStatePreempted: + // A preemption has already been delivered. + return nil + default: + panic("inconsistent state") + } + + ok := r.state.cas(state, target) + if ok && target == idleStatePreempted { + // We have entered preemption mode. + // Preempt the current read on the connection. + return r.conn.SetReadDeadline(time.Now()) + } + } +} + +// Read from the connection. +func (r *idleReader) Read(data []byte) (int, error) { + var needsDeadlineReset bool + state := r.state.load() + switch state { + case idleStateActive, idleStatePendingPreemption: + // Connection is active. + // There is no need to worry about preemption. + if r.timeout != 0 { + // Apply a read timeout. + err := r.conn.SetReadDeadline(time.Now().Add(r.timeout)) + if err != nil { + return 0, err + } + } + return r.conn.Read(data) + + case idleStateIdle: + // Read, and handle preemption. + n, err := r.conn.Read(data) + if err != nil { + // Check if the error was caused by preemption. + state = r.state.load() + switch { + case state == idleStatePreempted && n != 0: + // Some data was read before the preemption was delivered. + // Re-activate and discard the error. + + // Synchronize against the preempter. + // This is necessary to ensure that the cancellation deadline is cleared. + r.preemptMu.Lock() + defer r.preemptMu.Unlock() + + // Re-activate the connection. + // No CAS loop is necessary since we are synchronized against preempters. + r.state = idleStatePendingPreemption + + // The deadline may need to reset since the preempter may have changed it. + needsDeadlineReset = true + + case state == idleStatePreempted: + // The read was preempted. + return 0, errPreempted + + case state != idleStateIdle: + // No other states make sense here. + panic("inconsistent state") + + default: + // No preemption was involved. + // It is just a regular network error. + return n, err + } + } else { + // The read went through. + + // Exit from idle mode. + // Ideally, transition to active mode. + // However, a preemption may trigger while this is running. + for state == idleStateIdle { + if r.state.cas(idleStateIdle, idleStateActive) { + state = idleStateActive + break + } + + state = r.state.load() + } + + switch state { + case idleStateActive: + // The connection was reactivated normally. + + case idleStatePreempted: + // The connection was preempted after the read completed. + // Defer the preemption and complete successfully. + + // Synchronize against the preempter. + // This is necessary to ensure that the cancellation deadline is cleared. + r.preemptMu.Lock() + defer r.preemptMu.Unlock() + + // Re-activate the connection. + // No CAS loop is necessary since we are synchronized against preempters. + r.state = idleStatePendingPreemption + + // The deadline may need to reset since the preempter may have changed it. + needsDeadlineReset = true + + default: + panic("inconsistent state") + } + } + + if needsDeadlineReset && r.timeout == 0 { + // Clear the deadline. + err := r.conn.SetReadDeadline(time.Time{}) + if err != nil { + return n, err + } + } + + return n, nil + + case idleStatePreempted: + // The connection is preempted. + return 0, errPreempted + + default: + panic("inconsistent state") + } +} diff --git a/pg/message/io.go b/pg/message/io.go new file mode 100644 index 000000000..e09150b9d --- /dev/null +++ b/pg/message/io.go @@ -0,0 +1,124 @@ +// Copyright 2020 Pilosa Corp. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package message + +import ( + "bufio" + "encoding/binary" + "errors" + "io" +) + +// Reader reads messages. +type Reader interface { + ReadMessage() (Message, error) +} + +// Writer writes messages. +type Writer interface { + WriteMessage(Message) error + Flush() error +} + +// WireReader reads messages in Postgres wire protocol format. +type WireReader struct { + buf []byte + r *bufio.Reader + scratch [4]byte +} + +// ReadMessage reads a single message off of the wire. +// The returned message is only valid until the next read call, as the data buffer may be re-used. +func (r *WireReader) ReadMessage() (Message, error) { + t, err := r.r.ReadByte() + if err != nil { + return Message{}, err + } + + _, err = r.r.Read(r.scratch[:]) + if err != nil { + return Message{}, err + } + + len := binary.BigEndian.Uint32(r.scratch[:4]) + if len < 4 { + return Message{}, errors.New("invalid message length") + } + len -= 4 + + if cap(r.buf) < int(len) { + r.buf = make([]byte, len) + } else { + r.buf = r.buf[:len] + } + _, err = io.ReadFull(r.r, r.buf) + if err != nil { + return Message{}, err + } + + return Message{ + Type: Type(t), + Data: r.buf, + }, nil +} + +var _ Reader = (*WireReader)(nil) + +// NewWireReader returns a message reader that reads postgres wire protocol format. +func NewWireReader(r *bufio.Reader) *WireReader { + return &WireReader{r: r} +} + +// ErrMessageTooBig is an error indicating that a message is too big to be sent or received. +var ErrMessageTooBig = errors.New("message is too big") + +// WireWriter writes messages in Postgres wire protocol. +type WireWriter struct { + w *bufio.Writer + scratch [4]byte +} + +// WriteMessage writes a message onto the wire. +func (w *WireWriter) WriteMessage(message Message) error { + if uint(len(message.Data))+4 >= 1<<31 { + return ErrMessageTooBig + } + + err := w.w.WriteByte(byte(message.Type)) + if err != nil { + return err + } + + binary.BigEndian.PutUint32(w.scratch[:], uint32(len(message.Data))+4) + _, err = w.w.Write(w.scratch[:]) + if err != nil { + return err + } + + _, err = w.w.Write(message.Data) + return err +} + +// Flush writes any buffered data to the underlying stream. +func (w *WireWriter) Flush() error { + return w.w.Flush() +} + +var _ Writer = (*WireWriter)(nil) + +// NewWireWriter returns a message writer that writes in postgres wire protocol format. +func NewWireWriter(w *bufio.Writer) *WireWriter { + return &WireWriter{w: w} +} diff --git a/pg/message/message.go b/pg/message/message.go new file mode 100644 index 000000000..b7de6683f --- /dev/null +++ b/pg/message/message.go @@ -0,0 +1,344 @@ +// Copyright 2020 Pilosa Corp. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package message + +import ( + "bytes" + "encoding/binary" +) + +// Type is a byte indicating the type of a Postgres message. +type Type byte + +const ( + // TypeAuthentication is a message used to transfer authentication info. + TypeAuthentication Type = 'R' + + // TypeReadyForQuery is a message used to indicate that the server is ready for another query. + TypeReadyForQuery Type = 'Z' + + // TypeCommandComplete is a message used to indicate that a query has completed. + TypeCommandComplete Type = 'C' + + // TypeError is an error message. + TypeError Type = 'E' + + // TypeRowDescription is a message indicating the column types of the result rows from a query. + TypeRowDescription Type = 'T' + + // TypeDataRow is a message with the contents of a single row. + TypeDataRow Type = 'D' + + // TypeTermination is a message indicating a request to terminate a connection. + TypeTermination Type = 'X' + + // TypeNegotiateProtocolVersion is a message used when a client attempts to connect with a newer minor version than the server supports. + TypeNegotiateProtocolVersion Type = 'v' + + // TypeSimpleQuery is a simple query request. + TypeSimpleQuery Type = 'Q' +) + +// AuthenticationOK is a message indicating that authentication has completed. +var AuthenticationOK = Message{ + Type: TypeAuthentication, + Data: []byte{0, 0, 0, 0}, +} + +// Message is a Postgres message value. +type Message struct { + Type Type + Data []byte +} + +// TransactionStatus is the current transaction state. +type TransactionStatus byte + +const ( + // TransactionStatusIdle indicates that there is no active transaction. + TransactionStatusIdle TransactionStatus = 'I' + + // TransactionStatusActive indicates that the connection currently has an active transaction. + TransactionStatusActive TransactionStatus = 'T' + + // TransactionStatusFailed indicates that the connection currently has a failed transaction. + TransactionStatusFailed TransactionStatus = 'E' +) + +// Encoder encodes messages. +type Encoder struct { + buf bytes.Buffer + scratch [4]byte +} + +func (e *Encoder) i16(i int16) error { + binary.BigEndian.PutUint16(e.scratch[:2], uint16(i)) + _, err := e.buf.Write(e.scratch[:2]) + return err +} + +func (e *Encoder) i32(i int32) error { + binary.BigEndian.PutUint32(e.scratch[:], uint32(i)) + _, err := e.buf.Write(e.scratch[:]) + return err +} + +// ReadyForQuery encodes a "ready for query" message. +func (e *Encoder) ReadyForQuery(status TransactionStatus) (Message, error) { + e.buf.Reset() + + err := e.buf.WriteByte(byte(status)) + if err != nil { + return Message{}, err + } + + return Message{ + Type: TypeReadyForQuery, + Data: e.buf.Bytes(), + }, nil +} + +// CommandComplete encodes a command completion message. +func (e *Encoder) CommandComplete(tag string) (Message, error) { + e.buf.Reset() + + _, err := e.buf.WriteString(tag) + if err != nil { + return Message{}, err + } + + err = e.buf.WriteByte(0) + if err != nil { + return Message{}, err + } + + return Message{ + Type: TypeCommandComplete, + Data: e.buf.Bytes(), + }, nil +} + +// NoticeFieldType indicates the type of a notice/error field. +// https://www.postgresql.org/docs/9.3/protocol-error-fields.html +type NoticeFieldType byte + +const ( + // NoticeFieldSeverity indicates the severity of a notice/error. + NoticeFieldSeverity NoticeFieldType = 'S' + + // NoticeFieldMessage is a short human-readable error/notice message. + NoticeFieldMessage NoticeFieldType = 'M' + + // NoticeFieldDetail is an optional extended description of the error. + NoticeFieldDetail NoticeFieldType = 'D' + + // NoticeFieldHint is a suggestion of how to address the issue. + NoticeFieldHint NoticeFieldType = 'H' +) + +// NoticeField is a field in an error or notice. +type NoticeField struct { + Type NoticeFieldType + Data string +} + +func (e *Encoder) messageOrNotice(fields ...NoticeField) error { + for _, f := range fields { + err := e.buf.WriteByte(byte(f.Type)) + if err != nil { + return err + } + + _, err = e.buf.WriteString(f.Data) + if err != nil { + return err + } + + err = e.buf.WriteByte(0) + if err != nil { + return err + } + } + + return e.buf.WriteByte(0) +} + +// Error encodes a Postgres error message. +func (e *Encoder) Error(fields ...NoticeField) (Message, error) { + e.buf.Reset() + err := e.messageOrNotice(fields...) + if err != nil { + return Message{}, err + } + return Message{ + Type: TypeError, + Data: e.buf.Bytes(), + }, nil +} + +// GoError creates a simple Postgres error message from a Go error value. +func (e *Encoder) GoError(err error) (Message, error) { + return e.Error( + NoticeField{ + Type: NoticeFieldSeverity, + Data: "ERROR", + }, + NoticeField{ + Type: NoticeFieldMessage, + Data: err.Error(), + }, + ) +} + +// ColumnDescription is a description of a data column. +type ColumnDescription struct { + Name string + TableID int32 //either a table/col id or 0 + FieldID int16 //either a table/col id or 0 + TypeID int32 //field type + TypeLen int16 //size in bytes of field + TypeModifier int32 //type modifer? + Mode int16 //0=text 1=binary +} + +// RowDescription describes the response rows from a query. +func (e *Encoder) RowDescription(cols ...ColumnDescription) (Message, error) { + if len(cols) >= 1<<15 { + return Message{}, ErrMessageTooBig + } + + e.buf.Reset() + + err := e.i16(int16(len(cols))) + if err != nil { + return Message{}, nil + } + + for _, col := range cols { + _, err := e.buf.WriteString(col.Name) + if err != nil { + return Message{}, err + } + err = e.buf.WriteByte(0) + if err != nil { + return Message{}, err + } + + err = e.i32(col.TableID) + if err != nil { + return Message{}, err + } + + err = e.i16(col.FieldID) + if err != nil { + return Message{}, err + } + + err = e.i32(col.TypeID) + if err != nil { + return Message{}, err + } + + err = e.i16(col.TypeLen) + if err != nil { + return Message{}, err + } + + err = e.i32(col.TypeModifier) + if err != nil { + return Message{}, err + } + + err = e.i16(col.Mode) + if err != nil { + return Message{}, err + } + } + + return Message{ + Type: TypeRowDescription, + Data: e.buf.Bytes(), + }, nil +} + +// TextRow encodes a data row in textual format. +func (e *Encoder) TextRow(row ...string) (Message, error) { + if len(row) >= 1<<15 { + return Message{}, ErrMessageTooBig + } + + e.buf.Reset() + + err := e.i16(int16(len(row))) + if err != nil { + return Message{}, err + } + + for _, val := range row { + if uint(len(val)) >= 1<<31 { + return Message{}, ErrMessageTooBig + } + + err = e.i32(int32(len(val))) + if err != nil { + return Message{}, err + } + + _, err = e.buf.WriteString(val) + if err != nil { + return Message{}, err + } + } + + return Message{ + Type: TypeDataRow, + Data: e.buf.Bytes(), + }, nil +} + +// NegotiateProtocolVersion encodes a protocol negotiation packet. +func (e *Encoder) NegotiateProtocolVersion(maxMinor int32, unrecognizedOptions ...string) (Message, error) { + if uint64(len(unrecognizedOptions)) >= 1<<31 { + return Message{}, ErrMessageTooBig + } + + e.buf.Reset() + + err := e.i32(maxMinor) + if err != nil { + return Message{}, err + } + + err = e.i32(int32(len(unrecognizedOptions))) + if err != nil { + return Message{}, err + } + for _, opt := range unrecognizedOptions { + _, err = e.buf.WriteString(opt) + if err != nil { + return Message{}, err + } + + err = e.buf.WriteByte(0) + if err != nil { + return Message{}, err + } + } + + return Message{ + Type: TypeNegotiateProtocolVersion, + Data: e.buf.Bytes(), + }, nil +} diff --git a/pg/pgtest/handler.go b/pg/pgtest/handler.go new file mode 100644 index 000000000..e6cb22e12 --- /dev/null +++ b/pg/pgtest/handler.go @@ -0,0 +1,31 @@ +// Copyright 2020 Pilosa Corp. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pgtest + +import ( + "context" + + "github.com/pilosa/pilosa/v2/pg" +) + +// HandlerFunc implements a postgres query handler with a function. +type HandlerFunc func(context.Context, pg.QueryResultWriter, pg.Query) error + +// HandleQuery calls the user's query handler function. +func (h HandlerFunc) HandleQuery(ctx context.Context, w pg.QueryResultWriter, q pg.Query) error { + return h(ctx, w, q) +} + +var _ pg.QueryHandler = HandlerFunc(nil) diff --git a/pg/pgtest/memnet.go b/pg/pgtest/memnet.go new file mode 100644 index 000000000..0ba85565c --- /dev/null +++ b/pg/pgtest/memnet.go @@ -0,0 +1,71 @@ +// Copyright 2020 Pilosa Corp. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pgtest + +import ( + "errors" + "net" + "sync" +) + +// errListenerClosed is an error returned when the listener is closed. +var errListenerClosed = errors.New("listener closed") + +type inMemoryListener struct { + ch chan net.Conn + closed chan struct{} + once sync.Once +} + +func (l *inMemoryListener) Accept() (net.Conn, error) { + select { + case <-l.closed: + return nil, errListenerClosed + default: + } + select { + case conn := <-l.ch: + return conn, nil + case <-l.closed: + return nil, errListenerClosed + } +} + +func (l *inMemoryListener) Close() error { + l.once.Do(func() { close(l.closed) }) + + return nil +} + +type memAddr struct{} + +func (a memAddr) Network() string { return "memory" } +func (a memAddr) String() string { return "memory" } + +func (l *inMemoryListener) Addr() net.Addr { + return memAddr{} +} + +func (l *inMemoryListener) Dial() (net.Conn, error) { + serverConn, clientConn := net.Pipe() + select { + case l.ch <- serverConn: + return clientConn, nil + case <-l.closed: + serverConn.Close() + clientConn.Close() + return nil, errListenerClosed + } +} diff --git a/pg/pgtest/server.go b/pg/pgtest/server.go new file mode 100644 index 000000000..e3a48e244 --- /dev/null +++ b/pg/pgtest/server.go @@ -0,0 +1,91 @@ +// Copyright 2020 Pilosa Corp. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pgtest + +import ( + "context" + "net" + "testing" + + "github.com/pilosa/pilosa/v2/pg" + "github.com/pkg/errors" + "golang.org/x/sync/errgroup" +) + +// ShutdownFunc is a function to use to shut down a test fixture. +// This function will send a shutdown signal and then wait for completion. +type ShutdownFunc func() error + +// Finish invokes the shutdown function and fails the test if an error occurs. +func (f ShutdownFunc) Finish(tb testing.TB, name string) { + err := f() + if err != nil { + tb.Errorf("failed to shut down %s: %v", name, err) + } +} + +// ServeTCP creates a TCP listener and serves postgres wire protocol on it. +func ServeTCP(addr string, server *pg.Server) (net.Addr, ShutdownFunc, error) { + listener, err := net.Listen("tcp", addr) + if err != nil { + return nil, nil, errors.Wrap(err, "listening on TCP") + } + + laddr := listener.Addr() + + ctx, cancel := context.WithCancel(context.Background()) + var eg errgroup.Group + eg.Go(func() error { return server.Serve(ctx, listener) }) + + return laddr, + func() error { + cancel() + return eg.Wait() + }, + nil +} + +// ServeTLS sets up TLS on the server and invokes ServeTCP. +func ServeTLS(addr string, server *pg.Server) (net.Addr, ShutdownFunc, error) { + err := SetupTLS(server) + if err != nil { + return nil, nil, errors.Wrap(err, "server TLS setup failed") + } + + return ServeTCP(addr, server) +} + +// ConnectFunc is a function to connect to a server. +type ConnectFunc func() (net.Conn, error) + +// ServeMem serves postgres on in-memory connections. +// TLS does not work here, as it relies on the OS to buffer and discard data. +func ServeMem(server *pg.Server) (ConnectFunc, ShutdownFunc, error) { + listener := &inMemoryListener{ + ch: make(chan net.Conn), + closed: make(chan struct{}), + } + + ctx, cancel := context.WithCancel(context.Background()) + var eg errgroup.Group + eg.Go(func() error { return server.Serve(ctx, listener) }) + + return listener.Dial, + func() error { + cancel() + return eg.Wait() + }, + nil +} diff --git a/pg/pgtest/tls.go b/pg/pgtest/tls.go new file mode 100644 index 000000000..53e942e80 --- /dev/null +++ b/pg/pgtest/tls.go @@ -0,0 +1,101 @@ +// Copyright 2020 Pilosa Corp. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pgtest + +import ( + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "encoding/pem" + "math/big" + "time" + + "github.com/pilosa/pilosa/v2/pg" + "github.com/pkg/errors" +) + +// SetupTLS generates a TLS certificate and installs it into the server. +// TODO: have the client properly trust this (generate a CA to install instead of using self-signed). +func SetupTLS(server *pg.Server) error { + // Generate an ecdsa key for the cert. + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + return errors.Wrap(err, "generating TLS key") + } + + // Generate a random 128-bit serial number. + serialNumber, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128)) + if err != nil { + return errors.Wrap(err, "generating serial number") + } + + // Make the certificate valid starting now. + now := time.Now() + + // Create a certificate template. + template := x509.Certificate{ + SerialNumber: serialNumber, + Subject: pkix.Name{ + Organization: []string{"Molecula"}, + }, + NotBefore: now, + NotAfter: now.Add(time.Hour), + + KeyUsage: x509.KeyUsageDigitalSignature, + ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}, + BasicConstraintsValid: true, + } + + // Generate a self-signed x509 cert from the template and the key. + certData, err := x509.CreateCertificate(rand.Reader, &template, &template, key.Public(), key) + if err != nil { + return errors.Wrap(err, "encoding certificate x509") + } + + // Encode the cert to PEM so that the TLS package can load it. + certPEM := pem.EncodeToMemory(&pem.Block{ + Type: "CERTIFICATE", + Bytes: certData, + }) + + // Encode the private key to x509. + keyData, err := x509.MarshalPKCS8PrivateKey(key) + if err != nil { + return errors.Wrap(err, "encoding key x509") + } + + // Encode the key to PEM so that the TLS package can load it. + keyPEM := pem.EncodeToMemory(&pem.Block{ + Type: "PRIVATE KEY", + Bytes: keyData, + }) + + // Load the certificate and key from their PEM encodings. + cert, err := tls.X509KeyPair(certPEM, keyPEM) + if err != nil { + return errors.Wrap(err, "loading TLS key pair") + } + + // Install the certificate into the server. + if server.TLSConfig == nil { + server.TLSConfig = &tls.Config{} + } + server.TLSConfig.Certificates = append(server.TLSConfig.Certificates, cert) + + return nil +} diff --git a/pg/protocol.go b/pg/protocol.go new file mode 100644 index 000000000..57dc2168f --- /dev/null +++ b/pg/protocol.go @@ -0,0 +1,471 @@ +// Copyright 2020 Pilosa Corp. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pg + +import ( + "bufio" + "bytes" + "context" + "crypto/tls" + "encoding/binary" + "encoding/hex" + "fmt" + "io" + "io/ioutil" + "net" + "strings" + "sync" + "time" + + "github.com/pilosa/pilosa/v2/pg/message" + "github.com/pkg/errors" +) + +// Protocol is a Postgres protocol version. +type Protocol uint32 + +const ( + // ProtocolPostgres30 is version 3.0 of the Postgres wire protocol. + ProtocolPostgres30 Protocol = (3 << 16) | 0 + + // ProtocolCancel is the protocol used for query cancellation. + ProtocolCancel Protocol = (1234 << 16) | 5678 + + // ProtocolSSL is the protocol used for SSL upgrades. + ProtocolSSL Protocol = (1234 << 16) | 5679 + + // ProtocolSupported is the main protocol version supported by this package. + ProtocolSupported Protocol = ProtocolPostgres30 +) + +// Major returns the major revision of the protocol. +func (p Protocol) Major() uint16 { + return uint16(p >> 16) +} + +// Minor returns the minor revision of the protocol. +func (p Protocol) Minor() uint16 { + return uint16(p) +} + +func (p Protocol) String() string { + switch p { + case ProtocolCancel: + return "cancel" + case ProtocolSSL: + return "SSL" + } + + return fmt.Sprintf("v%d.%d", p.Major(), p.Minor()) +} + +// handle reads the startup packet and dispatches an appropriate protocol handler for the connection. +func (s *Server) handle(ctx context.Context, conn net.Conn) (err error) { + defer func() { + cerr := conn.Close() + if cerr != nil && err == nil { + err = errors.Wrap(cerr, "closing connection") + } + }() + + if tcpconn, ok := conn.(*net.TCPConn); ok { + // Postgres does not have any real mechanism for confirming that a connection is still alive. + // Without this, a connection that breaks while idle would live indefinitely. + // With a TCP keepalive, this should return an error after approximately 2 hours (depending on OS configuration). + err := tcpconn.SetKeepAlive(true) + if err != nil { + return errors.Wrap(err, "enabling TCP keepalive") + } + } + + var startupDeadline time.Time + if s.StartupTimeout > 0 { + // Set deadline for processing the startup. + startupDeadline = time.Now().Add(s.StartupTimeout) + err = conn.SetDeadline(startupDeadline) + if err != nil { + return errors.Wrap(err, "setting deadline on protocol startup") + } + } + +startup: + // Read startup packet. + var buf [4]byte + _, err = io.ReadFull(conn, buf[:]) + if err != nil { + return errors.Wrap(err, "reading startup message length") + } + size := binary.BigEndian.Uint32(buf[:]) + if size < 4 { + return errors.Errorf("invalid startup packet length: %d bytes", size) + } + maxLen := s.MaxStartupSize + if maxLen == 0 { + maxLen = 1024 * 1024 + } + if size > maxLen { + return errors.Errorf("oversized startup frame of %d bytes (max: %d bytes)", size, maxLen) + } + data := make([]byte, size-4) + _, err = io.ReadFull(conn, data) + if err != nil { + return errors.Wrap(err, "reading startup packet") + } + + // Extract protocol ID. + if len(data) < 4 { + return errors.Errorf("startup packet is too small for protocol ID: %d bytes", len(data)) + } + proto := Protocol(binary.BigEndian.Uint32(data)) + data = data[4:] + + // Handle special protocols. + switch proto { + case ProtocolCancel: + // TODO: send an actual postgres error message. + return errors.New("cancellation protocol not yet supported") + + case ProtocolSSL: + if s.TLSConfig != nil { + // Upgrade the connection to TLS and renegotiate on the tunneled connection. + _, err = conn.Write([]byte{'S'}) + if err != nil { + return errors.Wrap(err, "sending SSL support confirmation") + } + conn = tls.Server(conn, s.TLSConfig) + if s.StartupTimeout > 0 { + err := conn.SetDeadline(startupDeadline) + if err != nil { + return errors.Wrap(err, "transferring startup deadline to TLS connection") + } + } + goto startup + } + + // Inform the client that SSL is not available and try again. + s.Logger.Debugf("client at %s requested a secure postgres connection but TLS is not configured", conn.RemoteAddr()) + _, err = conn.Write([]byte{'N'}) + if err != nil { + return errors.Wrap(err, "sending SSL unsupported notification") + } + goto startup + } + + // Handle regular postgres. + return s.handleStandard(ctx, proto, conn, data) +} + +// parseParams parses a parameter list from a startup packet. +func parseParams(data []byte) (map[string]string, error) { + params := make(map[string]string) + for { + idx := bytes.IndexByte(data, 0) + switch idx { + case 0: + return params, nil + case -1: + return nil, errors.New("malformed startup parameter list") + } + + key := string(data[:idx]) + data = data[idx+1:] + + idx = bytes.IndexByte(data, 0) + if idx == -1 { + return nil, errors.New("malformed startup parameter list") + } + val := string(data[:idx]) + data = data[idx+1:] + + params[key] = val + } +} + +// handleStandard handles a connection in the standard postgres wire protocol. +// The client is responsible for closing the connection when this finishes. +func (s *Server) handleStandard(ctx context.Context, proto Protocol, conn net.Conn, data []byte) error { + // Wait for helper goroutines to finish. + var wg sync.WaitGroup + defer wg.Wait() + + // Set up context. + ctx, cancel := context.WithCancel(ctx) + defer cancel() + + // Check the major version. + if proto.Major() != ProtocolSupported.Major() { + return errors.Errorf("unsupported protocol %v", proto) + } + + // Parse the parameters bundled in the startup packet. + params, err := parseParams(data) + if err != nil { + return errors.Wrap(err, "parsing parameters") + } + + if user, ok := params["user"]; ok { + // Log the connection. + s.Logger.Debugf("new postgres connection from user %q at %v", user, conn.RemoteAddr()) + } else { + // We do not use this much yet, but the wire protocol says that it is required. + return errors.New("missing username") + } + + // Set up message input and output. + + // Set up a reader that will preempt the connection when the context is canceled. + ir := idleReader{ + conn: conn, + timeout: s.ReadTimeout, + } + wg.Add(1) + go func() { + defer wg.Done() + + <-ctx.Done() + + ir.preempt() //nolint:errcheck + }() + + // Clear the startup deadline. + err = conn.SetDeadline(time.Time{}) + if err != nil { + return err + } + + // Set up a message reader with buffering. + rbuf := bufio.NewReader(&ir) + r := message.NewWireReader(rbuf) + + // Set up a writer on the connection. + var ww io.Writer = conn + if s.WriteTimeout != 0 { + // Apply the write timeout. + ww = &timeoutWriter{ + conn: conn, + timeout: s.WriteTimeout, + } + } + + // Set up a message writer with buffering. + w := message.NewWireWriter(bufio.NewWriter(ww)) + + var encoder message.Encoder + if proto.Minor() > ProtocolSupported.Minor() { + // Negotiate the version down. + s.Logger.Debugf("client requested unsupported protocol version %v; attempting to downgrade to %v", proto, ProtocolSupported) + msg, err := encoder.NegotiateProtocolVersion(int32(ProtocolSupported.Minor())) + if err != nil { + return errors.Wrap(err, "negotiating version") + } + err = w.WriteMessage(msg) + if err != nil { + return errors.Wrap(err, "negotiating version") + } + } + + // TODO: real auth + err = w.WriteMessage(message.AuthenticationOK) + if err != nil { + return errors.Wrap(err, "sending authentication confirmation") + } + + var queryReady bool + for { + if !queryReady { + // Indicate that we are ready for a query. + // TODO: provide a valid transaction state. + msg, err := encoder.ReadyForQuery(message.TransactionStatusActive) + if err != nil { + return errors.Wrap(err, "sending query ready status") + } + err = w.WriteMessage(msg) + if err != nil { + return errors.Wrap(err, "sending query ready status") + } + + // Flush the write buffer so that the client can respond. + err = w.Flush() + if err != nil { + return errors.Wrap(err, "flushing status") + } + + if rbuf.Buffered() == 0 { + // Put the connection into idle mode. + err = ir.setIdle() + if err != nil { + return errors.Wrap(err, "setting idle mode") + } + } else { + // If the client follows the spec, then it should not have sent anything more. + // However, it seems that no clients completely follow the spec, so we shouldn't rely on anything that isn't entirely straightforward. + s.Logger.Debugf("postgres client sent additional data without waiting for completion") + } + } + + // Read the next packet. + msg, err := r.ReadMessage() + if err != nil { + if err == errPreempted { + // The server is shutting down. + return errors.Wrap(s.handleShutdown( + conn, w, &encoder, + + message.NoticeField{ + Type: message.NoticeFieldSeverity, + Data: "ERROR", + }, + message.NoticeField{ + Type: message.NoticeFieldMessage, + Data: "server shutting down", + }, + message.NoticeField{ + Type: message.NoticeFieldHint, + Data: "This is normal. This message is sent when a server is shutting down and terminating its connections.", + }, + ), "processing connection shutdown") + } + + return err + } + + switch msg.Type { + case message.TypeTermination: + // We are done. + return w.Flush() + + case message.TypeSimpleQuery: + // Execute a simple query. + + queryReady = false + + // Parse the query message (a null-terminated string). + query := SimpleQuery(strings.TrimSuffix(string(msg.Data), "\x00")) + + // Set up a result writer. + // SELECT is used as a default tag, which seems to be handled decently by clients. + // The encoder is intentionally not used because its buffer may be huge. + qwriter := &queryResultWriter{ + w: w, + te: s.TypeEngine, + tag: "SELECT", + } + + // Dispatch the query handler. + // TODO: cancellation (requires crazy internode logic and fake process IDs) + // This is not the connection context, since we want the request to finish safely before connection shutdown. + qerr := s.QueryHandler.HandleQuery(context.Background(), qwriter, query) + if qerr != nil { + // There was an error in processing the query. + // Send the error back to the client and keep going. + s.Logger.Debugf("failed to execute query %q: %v", query, qerr) + msg, err = encoder.GoError(qerr) + if err != nil { + return errors.Wrap(err, "failed to send query error to client") + } + err = w.WriteMessage(msg) + if err != nil { + return errors.Wrap(err, "failed to send query error to client") + } + } else { + // The query completed normally. + // Notify the client of completion. + msg, err = encoder.CommandComplete(qwriter.tag) + if err != nil { + return errors.Wrap(err, "sending command completion notification") + } + err = w.WriteMessage(msg) + if err != nil { + return errors.Wrap(err, "sending command completion notification") + } + } + + // The data will be flushed after we write back the "ready for query" state. + + default: + // The message is not supported yet. + // Send an error. + s.Logger.Printf("unrecognized postgres packet %v", msg) + msg, err = encoder.Error( + message.NoticeField{ + Type: message.NoticeFieldSeverity, + Data: "ERROR", + }, + message.NoticeField{ + Type: message.NoticeFieldMessage, + Data: fmt.Sprintf("unrecognized message type %q", msg.Type), + }, + message.NoticeField{ + Type: message.NoticeFieldDetail, + Data: "message body:" + hex.Dump(msg.Data), + }, + ) + if err != nil { + return errors.Wrap(err, "sending unrecognized message error") + } + err = w.WriteMessage(msg) + if err != nil { + return errors.Wrap(err, "sending unrecognized message error") + } + err = w.Flush() + if err != nil { + return errors.Wrap(err, "sending unrecognized message error") + } + } + } +} + +func (s *Server) handleShutdown(conn net.Conn, w message.Writer, encoder *message.Encoder, notice ...message.NoticeField) error { + var wg sync.WaitGroup + defer wg.Wait() + + // Try to send a message to the client before closing the connection. + msg, err := encoder.Error(notice...) + if err != nil { + return errors.Wrap(err, "generating shutdown notification") + } + + if s.WriteTimeout == 0 { + // The client is likely to not listen for incoming messages. + // Force a write timeout to ensure that this terminates. + err := conn.SetWriteDeadline(time.Now().Add(time.Second)) + if err != nil { + return errors.Wrap(err, "setting shutdown write deadline") + } + } + + // The client may be waiting on a write, so we need to drain the incoming data stream. + err = conn.SetReadDeadline(time.Time{}) + if err != nil { + return errors.Wrap(err, "clearing read deadline for shutdown") + } + defer conn.SetReadDeadline(time.Now()) //nolint:errcheck + wg.Add(1) + go func() { + defer wg.Done() + + io.Copy(ioutil.Discard, conn) //nolint:errcheck + }() + + // Attempt to send the shutdown notification. + // This will fail under many scenarios, as the client is not necessarily reading. + err = w.WriteMessage(msg) + if err != nil { + return nil + } + w.Flush() + + return nil +} diff --git a/pg/query.go b/pg/query.go new file mode 100644 index 000000000..8b95b4df8 --- /dev/null +++ b/pg/query.go @@ -0,0 +1,130 @@ +// Copyright 2020 Pilosa Corp. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pg + +import ( + "context" + "fmt" + + "github.com/pilosa/pilosa/v2/pg/message" + "github.com/pkg/errors" +) + +// Query is an interface to be implemented by queries. +type Query interface { + fmt.Stringer +} + +// SimpleQuery is a query sent as only a string. +// It has no parameters. +type SimpleQuery string + +func (q SimpleQuery) String() string { + return string(q) +} + +// ColumnInfo contains metadata about a column. +type ColumnInfo struct { + Name string + Type Type + TableID int32 + FieldID int16 +} + +// QueryResultWriter is used to write the results of a query back over the connection. +type QueryResultWriter interface { + // WriteHeader sets the column header information. + WriteHeader(...ColumnInfo) error + + // WriteRowText sends a row of data in textual format. + WriteRowText(...string) error + + // Tag assigns a tag to the query. + // This should be called before the query is completed. + Tag(tag string) +} + +// QueryHandler handles a query. +type QueryHandler interface { + // HandleQuery executes a query and writes the results back. + HandleQuery(context.Context, QueryResultWriter, Query) error +} + +// queryResultWriter implements QueryResultWrtiter over postgres wire protocol. +// The underlying message writer must be flushed by the caller once the query has finished. +type queryResultWriter struct { + w message.Writer + te TypeEngine + enc message.Encoder + width int + wroteHeaders bool + tag string +} + +func (w *queryResultWriter) WriteHeader(info ...ColumnInfo) error { + if w.wroteHeaders { + return errors.New("double-write of query headers") + } + + // Translate column information into a row description message. + desc := make([]message.ColumnDescription, len(info)) + for i, c := range info { + t, err := w.te.TranslateType(c.Type) + if err != nil { + return errors.Wrap(err, "translating column type") + } + t.Name = c.Name + t.TableID = c.TableID + t.FieldID = c.FieldID + desc[i] = t + } + + // Encode the row description. + msg, err := w.enc.RowDescription(desc...) + if err != nil { + return errors.Wrap(err, "encoding query header") + } + + w.wroteHeaders = true + w.width = len(desc) + + // Write the row description. + return w.w.WriteMessage(msg) +} + +func (w *queryResultWriter) WriteRowText(text ...string) error { + // Check preconditions of the call. + switch { + case !w.wroteHeaders: + return errors.New("writing rows without headers") + case len(text) != w.width: + return errors.Errorf("expected %d columns but found %d", w.width, len(text)) + } + + // Encode the row data as text into a DataRow message. + msg, err := w.enc.TextRow(text...) + if err != nil { + return err + } + + // Write the data row over the network. + return w.w.WriteMessage(msg) +} + +func (w *queryResultWriter) Tag(tag string) { + w.tag = tag +} + +var _ QueryResultWriter = (*queryResultWriter)(nil) diff --git a/pg/server.go b/pg/server.go new file mode 100644 index 000000000..55d374005 --- /dev/null +++ b/pg/server.go @@ -0,0 +1,136 @@ +// Copyright 2020 Pilosa Corp. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pg + +import ( + "context" + "crypto/tls" + "net" + "sync" + "time" + + "github.com/pilosa/pilosa/v2/logger" +) + +// Server is a postgres wire protocol server. +type Server struct { + // QueryHandler will be used to serve query requests. + QueryHandler QueryHandler + + // TypeEngine is the type engine to use to result columns in query requests. + TypeEngine TypeEngine + + // TLSConfig is the TLS configuration to use to serve postgres TLS connections. + TLSConfig *tls.Config + + // StartupTimeout is the timeout to use for connection startup. + // If a connection fails to set up a protocol before this completes, it will be terminated. + StartupTimeout time.Duration + + // ReadTimeout is the timeout to apply for active reads (reads during the lifetime of a command). + // This timeout does not apply to an idling connection. + ReadTimeout time.Duration + + // WriteTimeout is the timeout to apply to network writes. + WriteTimeout time.Duration + + // MaxStartupSize is the maximum size of the startup packet (in bytes). + // This defaults to 2^31-1 bytes, which is the maximum size allowed by the protocol. + MaxStartupSize uint32 + + // ConnectionLimit is the maximum number of connections to allow at once. + ConnectionLimit uint16 + + // Logger is the logger to use for error conditions and state changes. + Logger logger.Logger +} + +// ServeConn serves a single connection. +func (s *Server) ServeConn(ctx context.Context, conn net.Conn) error { + return s.handle(ctx, conn) +} + +// Serve accepts postgres connections from a listener and processes them. +// If the context is cancelled, this will stop accepting requests and wait until all connections have terminated. +// No error will be returned if terminated by context cancellation. +// This will close the connection for the caller. +func (s *Server) Serve(ctx context.Context, l net.Listener) (err error) { + // Ignore errors triggered by a shutdown. + // Also propogate any error from terminating the listener. + var cerr error + defer func(ctx context.Context) { + if ctx.Err() == context.Canceled { + err = cerr + } + }(ctx) + + // Wait for the listener to be closed and all connections to shut down. + var wg sync.WaitGroup + defer wg.Wait() + + // Wrap the context to propogate a shutdown to the listeners and connection handlers. + ctx, cancel := context.WithCancel(ctx) + defer cancel() + + // Start a goroutine to shut down the listener when the context is canceled. + wg.Add(1) + go func() { + defer wg.Done() + + <-ctx.Done() + cerr = l.Close() + }() + + // Set up a semaphore for the connection limit. + var limit chan struct{} + done := ctx.Done() + if s.ConnectionLimit != 0 { + limit = make(chan struct{}, s.ConnectionLimit) + } + + for { + if limit != nil { + // Wait for connection limit. + if len(limit) == cap(limit) { + s.Logger.Printf("postgres connection limit reached") + } + select { + case limit <- struct{}{}: + case <-done: + return nil + } + } + + // Accept a connection. + conn, err := l.Accept() + if err != nil { + return err + } + + // Handle the connection in another goroutine. + wg.Add(1) + go func() { + defer wg.Done() + if limit != nil { + // Restore connection limit when done. + defer func() { <-limit }() + } + err := s.handle(ctx, conn) + if err != nil { + s.Logger.Printf("postgres connection terminated with error: %v", err) + } + }() + } +} diff --git a/pg/server_test.go b/pg/server_test.go new file mode 100644 index 000000000..4903fdb60 --- /dev/null +++ b/pg/server_test.go @@ -0,0 +1,218 @@ +// Copyright 2020 Pilosa Corp. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pg_test + +import ( + "context" + "fmt" + "net" + "os/exec" + "strconv" + "testing" + "time" + + "github.com/lib/pq" + "github.com/pilosa/pilosa/v2/logger" + "github.com/pilosa/pilosa/v2/pg" + "github.com/pilosa/pilosa/v2/pg/pgtest" +) + +// TestStartupTimeout tests that an incoming connection that does nothing times out and gets closed. +func TestStartupTimeout(t *testing.T) { + t.Parallel() + + connect, shutdown, err := pgtest.ServeMem(&pg.Server{ + StartupTimeout: time.Millisecond, + Logger: logger.NewLogfLogger(t), + }) + if err != nil { + t.Fatalf("starting in-memory postgres server: %v", err) + } + defer shutdown.Finish(t, "in-memory postgres server") + + conn, err := connect() + if err != nil { + t.Fatalf("failed to acquire connection: %v", err) + } + defer conn.Close() + + // The server isn't sending anything, so this should block until the connection dies. + conn.Read(make([]byte, 1024)) //nolint:errcheck +} + +// TestStartupInvalidLength tests that sending an HTTP GET request does not cause the server to allocate 1.2 GiB of memory. +func TestStartupInvalidLength(t *testing.T) { + t.Parallel() + + res := testing.Benchmark(func(b *testing.B) { + connect, shutdown, err := pgtest.ServeMem(&pg.Server{ + MaxStartupSize: 1024, + Logger: logger.NewLogfLogger(t), + }) + if err != nil { + t.Fatalf("starting in-memory postgres server: %v", err) + } + defer shutdown.Finish(t, "in-memory postgres server") + + b.ReportAllocs() + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + conn, err := connect() + if err != nil { + t.Fatalf("failed to acquire connection: %v", err) + } + + _, err = conn.Write([]byte("GET ")) + if err != nil { + t.Fatalf("failed to write invalid length: %v", err) + } + + // The server isn't sending anything, so this should block until the connection dies. + conn.Read(make([]byte, 1024)) //nolint:errcheck + + err = conn.Close() + if err != nil { + t.Fatalf("failed to close connection: %v", err) + } + } + }) + bpo := res.AllocedBytesPerOp() + t.Logf("allocated %d bytes per op", bpo) + if bpo > 1024*1024 { + t.Errorf("allocated too much memory: %d bytes/connection", bpo) + } +} + +// TestPQConnect tests connecting the Go SQL driver `pq` to this postgres server. +func TestPQConnect(t *testing.T) { + t.Parallel() + + server := &pg.Server{ + StartupTimeout: time.Second, + Logger: logger.NewLogfLogger(t), + } + addr, shutdown, err := pgtest.ServeTCP(":0", server) + if err != nil { + t.Fatalf("starting postgres server: %v", err) + } + defer shutdown.Finish(t, "postgres server") + + tcpAddr := addr.(*net.TCPAddr) + + connector, err := pq.NewConnector(fmt.Sprintf("user=molecula dbname=pilosa sslmode=disable host=%s port=%d", tcpAddr.IP, tcpAddr.Port)) + if err != nil { + t.Fatalf("failed to create connector: %v", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + conn, err := connector.Connect(ctx) + if err != nil { + t.Fatalf("failed to connect to postgres: %v", err) + } + defer pgtest.ShutdownFunc(conn.Close).Finish(t, "postgres TLS conn") +} + +// TestPQConnectSSL tests connecting the Go SQL driver `pq` to this postgres server, with SSL enabled. +func TestPQConnectSSL(t *testing.T) { + t.Parallel() + + server := &pg.Server{ + StartupTimeout: time.Second, + Logger: logger.NewLogfLogger(t), + } + addr, shutdown, err := pgtest.ServeTLS(":0", server) + if err != nil { + t.Fatalf("starting postgres server: %v", err) + } + defer shutdown.Finish(t, "postgres TLS server") + + tcpAddr := addr.(*net.TCPAddr) + + connector, err := pq.NewConnector(fmt.Sprintf("user=molecula dbname=pilosa sslmode=require host=%s port=%d", tcpAddr.IP, tcpAddr.Port)) + if err != nil { + t.Fatalf("failed to create connector: %v", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + conn, err := connector.Connect(ctx) + if err != nil { + t.Fatalf("failed to connect to postgres: %v", err) + } + defer pgtest.ShutdownFunc(conn.Close).Finish(t, "postgres TLS conn") +} + +// TestPSQLQuery tests sending a query from the `psql` command line tool. +func TestPSQLQuery(t *testing.T) { + // Check if psql is present. + // Skip this test if it is not. + _, err := exec.LookPath("psql") + if err != nil { + if err, ok := err.(*exec.Error); ok { + if err.Err == exec.ErrNotFound { + t.Skip("psql is not available") + } + } + t.Fatalf("searching for psql: %v", err) + } + + server := &pg.Server{ + QueryHandler: pgtest.HandlerFunc(func(ctx context.Context, w pg.QueryResultWriter, q pg.Query) error { + err := w.WriteHeader(pg.ColumnInfo{ + Name: "field", + Type: pg.TypeCharoid, + }) + if err != nil { + return err + } + + err = w.WriteRowText("h") + if err != nil { + return err + } + + err = w.WriteRowText("xyzzy") + if err != nil { + return err + } + + return nil + }), + TypeEngine: pg.PrimitiveTypeEngine{}, + StartupTimeout: time.Second, + Logger: logger.NewLogfLogger(t), + } + addr, shutdown, err := pgtest.ServeTCP(":0", server) + if err != nil { + t.Fatalf("starting postgres server: %v", err) + } + defer shutdown.Finish(t, "postgres server") + + tcpAddr := addr.(*net.TCPAddr) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + cmd := exec.CommandContext(ctx, "psql", "-h", tcpAddr.IP.String(), "-p", strconv.Itoa(tcpAddr.Port), "-c", "test query") + data, err := cmd.CombinedOutput() + if err != nil { + t.Fatalf("psql failed: %v", string(data)) + } +} diff --git a/pg/type.go b/pg/type.go new file mode 100644 index 000000000..7a89cc947 --- /dev/null +++ b/pg/type.go @@ -0,0 +1,54 @@ +// Copyright 2020 Pilosa Corp. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pg + +import "github.com/pilosa/pilosa/v2/pg/message" + +// Type represents a postgres type. +type Type struct { + // I am not entirely sure what should be in here long term. + // For now, I am just going to leave it like this. + id int32 +} + +// TypeCharoid is a postgres type for text. +var TypeCharoid = Type{id: 18} + +// TypeData is a type containing raw postgres wire type information. +type TypeData struct { + TypeID int32 + TypeLen int16 + TypeModifier int32 +} + +// TypeEngine is a system for managing types. +// This is necessary for compound types like arrays which need ID generation. +type TypeEngine interface { + // TranslateType populates a column description with type information. + TranslateType(Type) (message.ColumnDescription, error) +} + +// PrimitiveTypeEngine is a simple type engine that only works on primitive types. +type PrimitiveTypeEngine struct{} + +// TranslateType translates a type to a column description. +func (pte PrimitiveTypeEngine) TranslateType(t Type) (message.ColumnDescription, error) { + return message.ColumnDescription{ + TypeID: t.id, // just charoid for now; as far as I can tell most implementations do not really use this + TypeLen: -1, // vdsm had 4. . . but the spec says this should be negative + TypeModifier: -1, + Mode: 0, // send as text + }, nil +} diff --git a/server/config.go b/server/config.go index 2da11a365..7437d65ed 100644 --- a/server/config.go +++ b/server/config.go @@ -167,6 +167,25 @@ type Config struct { MutexFraction int `toml:"mutex-fraction"` } `toml:"profile"` + Postgres struct { + // Addr is the address to which to bind a postgres endpoint. + // If this is empty, no endpoint will be created. + Addr string `toml:"addr"` + // TLS configuration for postgres connections. + TLS TLSConfig `toml:"tls"` + + StartupTimeout toml.Duration `toml:"startup-timeout"` + ReadTimeout toml.Duration `toml:"read-timeout"` + WriteTimeout toml.Duration `toml:"write-timout"` + + MaxStartupSize uint32 `toml:"max-startup-size"` + + // ConnectionLimit is the maximum number of postgres connections to allow simultaneously. + // Setting this to 0 disables the limit. + // This mostly exists because other DBs seem to have it. + ConnectionLimit uint16 `toml:"max-connections"` + } `toml:"postgres"` + // Txsrc determines which Tx implementation the holder/Index will use; one // of the available transactional-storage engines. Choices are listed // in the string constants below. Should be one of @@ -235,6 +254,13 @@ func NewConfig() *Config { c.Profile.BlockRate = 10000000 // 1 sample per 10 ms c.Profile.MutexFraction = 100 // 1% sampling + // Postgres config (off by default). + c.Postgres.MaxStartupSize = 8 * 1024 * 1024 + c.Postgres.StartupTimeout = toml.Duration(5 * time.Second) + c.Postgres.ReadTimeout = toml.Duration(10 * time.Second) + c.Postgres.WriteTimeout = toml.Duration(10 * time.Second) + // we don't really need a connection limit + return c } diff --git a/server/pg.go b/server/pg.go new file mode 100644 index 000000000..98117c374 --- /dev/null +++ b/server/pg.go @@ -0,0 +1,399 @@ +// Copyright 2020 Pilosa Corp. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package server + +import ( + "context" + "crypto/tls" + "encoding/json" + "fmt" + "net" + "strconv" + "strings" + "time" + + "github.com/pilosa/pilosa/v2" + "github.com/pilosa/pilosa/v2/logger" + "github.com/pilosa/pilosa/v2/pg" + pb "github.com/pilosa/pilosa/v2/proto" + + "github.com/pkg/errors" + "golang.org/x/sync/errgroup" +) + +// PostgresServer provides a postgres endpoint on pilosa. +type PostgresServer struct { + api *pilosa.API + logger logger.Logger + eg errgroup.Group + s pg.Server + stop context.CancelFunc +} + +// NewPostgresServer creates a postgres server. +func NewPostgresServer(api *pilosa.API, logger logger.Logger, tls *tls.Config) *PostgresServer { + return &PostgresServer{ + api: api, + logger: logger, + s: pg.Server{ + QueryHandler: &queryDecodeHandler{ + child: &pilosaQueryHandler{ + api: api, + }, + }, + TypeEngine: pg.PrimitiveTypeEngine{}, + StartupTimeout: 5 * time.Second, + ReadTimeout: 10 * time.Second, + WriteTimeout: 10 * time.Second, + MaxStartupSize: 8 * 1024 * 1024, + Logger: logger, + TLSConfig: tls, + }, + } +} + +// Start a postgres endpoint at the specified address. +func (s *PostgresServer) Start(addr string) error { + l, err := net.Listen("tcp", addr) + if err != nil { + return errors.Wrap(err, "creating listener") + } + + s.logger.Printf("serving postgres wire protocol on %s", l.Addr()) + + ctx, cancel := context.WithCancel(context.Background()) + s.stop = cancel + + s.eg.Go(func() error { return s.s.Serve(ctx, l) }) + + return nil +} + +func (s *PostgresServer) Close() error { + if s == nil { + return nil + } + + if s.stop == nil { + return nil + } + s.stop() + + s.logger.Printf("waiting for postgres connections to shut down") + + return s.eg.Wait() +} + +type pgPQLQuery struct { + index string + query string +} + +func (q pgPQLQuery) String() string { + return fmt.Sprintf("[%s]%s", q.index, q.query) +} + +func pgDecodePQL(str string) (q pg.Query, err error) { + defer func() { + err = errors.Wrap(err, "not a valid PQL-over-postgres query") + }() + + if !strings.HasPrefix(str, "[") { + return nil, errors.New("missing index specification") + } + idx := strings.IndexRune(str, ']') + if idx == -1 { + return nil, errors.New("unclosed bracket in index specification") + } + + return pgPQLQuery{ + index: str[1:idx], + query: str[idx+1:], + }, nil +} + +type pilosaQueryHandler struct { + api *pilosa.API +} + +func pgWriteRow(w pg.QueryResultWriter, row *pilosa.Row) error { + err := w.WriteHeader(pg.ColumnInfo{ + Name: "_id", + Type: pg.TypeCharoid, + }) + if err != nil { + return errors.Wrap(err, "writing result header") + } + + if row.Keys != nil { + for _, k := range row.Keys { + err = w.WriteRowText(k) + if err != nil { + return errors.Wrap(err, "writing key") + } + } + } else { + for _, col := range row.Columns() { + err = w.WriteRowText(strconv.FormatUint(col, 10)) + if err != nil { + return errors.Wrap(err, "writing column ID") + } + } + } + + return nil +} + +func pgWriteRows(w pg.QueryResultWriter, rows pilosa.RowIdentifiers) error { + err := w.WriteHeader(pg.ColumnInfo{ + Name: rows.Field(), + Type: pg.TypeCharoid, + }) + if err != nil { + return errors.Wrap(err, "writing result header") + } + + if rows.Keys != nil { + for _, k := range rows.Keys { + err = w.WriteRowText(k) + if err != nil { + return errors.Wrap(err, "writing key") + } + } + } else { + for _, row := range rows.Rows { + err = w.WriteRowText(strconv.FormatUint(row, 10)) + if err != nil { + return errors.Wrap(err, "writing row ID") + } + } + } + + return nil +} + +func pgFormatVal(val interface{}) string { + switch val := val.(type) { + case bool: + return strconv.FormatBool(val) + case int64: + return strconv.FormatInt(val, 10) + case uint64: + return strconv.FormatUint(val, 10) + case string: + return val + default: + data, _ := json.Marshal(val) + return string(data) + } +} + +func pgWriteExtractedTable(w pg.QueryResultWriter, tbl pilosa.ExtractedTable) error { + headers := make([]pg.ColumnInfo, len(tbl.Fields)+1) + headers[0] = pg.ColumnInfo{ + Name: "_id", + Type: pg.TypeCharoid, + } + dataHeaders := headers[1:] + for i, f := range tbl.Fields { + dataHeaders[i] = pg.ColumnInfo{ + Name: f.Name, + Type: pg.TypeCharoid, + } + } + err := w.WriteHeader(headers...) + if err != nil { + return errors.Wrap(err, "writing result header") + } + + vals := make([]string, len(headers)) + dataVals := vals[1:] + for _, col := range tbl.Columns { + if col.Column.Keyed { + vals[0] = col.Column.Key + } else { + vals[0] = strconv.FormatUint(col.Column.ID, 10) + } + for i, v := range col.Rows { + dataVals[i] = pgFormatVal(v) + } + err = w.WriteRowText(vals...) + if err != nil { + return errors.Wrap(err, "writing result row") + } + } + + return nil +} + +func pgWriteGroupCount(w pg.QueryResultWriter, counts []pilosa.GroupCount) error { + if len(counts) == 0 { + // Not enough information is available to construct the header. + // This is a significant flaw in the data type. + return nil + } + + headers := make([]pg.ColumnInfo, len(counts[0].Group)+2) + for i, g := range counts[0].Group { + headers[i] = pg.ColumnInfo{ + Name: g.Field, + Type: pg.TypeCharoid, + } + } + headers[len(headers)-2] = pg.ColumnInfo{ + Name: "count", + Type: pg.TypeCharoid, + } + headers[len(headers)-1] = pg.ColumnInfo{ + Name: "sum", + Type: pg.TypeCharoid, + } + err := w.WriteHeader(headers...) + if err != nil { + return errors.Wrap(err, "writing result header") + } + + vals := make([]string, len(headers)) + for _, gc := range counts { + for j, g := range gc.Group { + var v string + switch { + case g.Value != nil: + v = strconv.FormatInt(*g.Value, 10) + case g.RowKey != "": + v = g.RowKey + default: + v = strconv.FormatUint(g.RowID, 10) + } + vals[j] = v + } + vals[len(vals)-2] = strconv.FormatUint(gc.Count, 10) + vals[len(vals)-1] = strconv.FormatInt(gc.Sum, 10) + + err := w.WriteRowText(vals...) + if err != nil { + return errors.Wrap(err, "writing group count result") + } + } + + return nil +} + +func pgWriteRowser(w pg.QueryResultWriter, result pb.ToRowser) error { + var data []string + return result.ToRows(func(row *pb.RowResponse) error { + if data == nil { + headers := make([]pg.ColumnInfo, len(row.Columns)) + for i, h := range row.Headers { + headers[i] = pg.ColumnInfo{ + Name: h.Name, + Type: pg.TypeCharoid, + } + } + err := w.WriteHeader(headers...) + if err != nil { + return errors.Wrap(err, "writing headers") + } + + data = make([]string, len(headers)) + } + + for i, col := range row.Columns { + var v string + switch col := col.ColumnVal.(type) { + case *pb.ColumnResponse_BoolVal: + v = strconv.FormatBool(col.BoolVal) + case *pb.ColumnResponse_DecimalVal: + v = col.DecimalVal.String() + case *pb.ColumnResponse_Float64Val: + v = strconv.FormatFloat(col.Float64Val, 'g', -1, 64) + case *pb.ColumnResponse_Int64Val: + v = strconv.FormatInt(col.Int64Val, 10) + case *pb.ColumnResponse_Uint64Val: + v = strconv.FormatUint(col.Uint64Val, 10) + case *pb.ColumnResponse_StringVal: + v = col.StringVal + case *pb.ColumnResponse_StringArrayVal: + data, _ := json.Marshal(col.StringArrayVal.Vals) + v = string(data) + case *pb.ColumnResponse_Uint64ArrayVal: + data, _ := json.Marshal(col.Uint64ArrayVal.Vals) + v = string(data) + default: + return errors.Errorf("unable to process value of type %T", col) + } + + data[i] = v + } + + return w.WriteRowText(data...) + }) +} + +func pgWriteResult(w pg.QueryResultWriter, result interface{}) error { + switch result := result.(type) { + case *pilosa.Row: + return pgWriteRow(w, result) + case pilosa.RowIdentifiers: + return pgWriteRows(w, result) + case pilosa.ExtractedTable: + return pgWriteExtractedTable(w, result) + case []pilosa.GroupCount: + return pgWriteGroupCount(w, result) + case pb.ToRowser: // we should avoid protobuf where we can... + return pgWriteRowser(w, result) + default: + return errors.Errorf("result type %T not yet supported", result) + } +} + +func (pqh *pilosaQueryHandler) HandleQuery(ctx context.Context, w pg.QueryResultWriter, q pg.Query) error { + switch q := q.(type) { + case pgPQLQuery: + resp, err := pqh.api.Query(ctx, &pilosa.QueryRequest{ + Index: q.index, + Query: q.query, + }) + if err != nil { + return errors.Wrap(err, "executing query") + } + if len(resp.Results) != 1 { + return errors.Errorf("expected 1 query result but found %d", len(resp.Results)) + } + return errors.Wrap(pgWriteResult(w, resp.Results[0]), "writing query result") + + default: + return errors.Errorf("query type %T not yet supported (query: %s)", q, q) + } +} + +type queryDecodeHandler struct { + child pg.QueryHandler +} + +func (qdh *queryDecodeHandler) HandleQuery(ctx context.Context, w pg.QueryResultWriter, q pg.Query) error { + switch qv := q.(type) { + case pg.SimpleQuery: + if strings.HasPrefix(string(qv), "[") { + pqlQuery, err := pgDecodePQL(strings.TrimSuffix(string(qv), ";")) + if err != nil { + return errors.Wrap(err, "decoding query") + } + q = pqlQuery + } + } + + return qdh.child.HandleQuery(ctx, w, q) +} diff --git a/server/server.go b/server/server.go index 2b1cea4bb..5578cf35f 100644 --- a/server/server.go +++ b/server/server.go @@ -87,6 +87,7 @@ type Command struct { listenURI *pilosa.URI tlsConfig *tls.Config closeTimeout time.Duration + pgserver *PostgresServer serverOptions []pilosa.ServerOption } @@ -171,6 +172,29 @@ func (m *Command) Start() (err error) { } }() + // Initialize postgres. + m.pgserver = nil + if m.Config.Postgres.Addr != "" { + var tlsConf *tls.Config + if m.Config.Postgres.TLS.CertificatePath != "" { + conf, err := GetTLSConfig(&m.Config.Postgres.TLS, m.logger.Logger()) + if err != nil { + return errors.Wrap(err, "settuing up postgres TLS") + } + tlsConf = conf + } + m.pgserver = NewPostgresServer(m.API, m.logger, tlsConf) + m.pgserver.s.StartupTimeout = time.Duration(m.Config.Postgres.StartupTimeout) + m.pgserver.s.ReadTimeout = time.Duration(m.Config.Postgres.ReadTimeout) + m.pgserver.s.WriteTimeout = time.Duration(m.Config.Postgres.WriteTimeout) + m.pgserver.s.MaxStartupSize = m.Config.Postgres.MaxStartupSize + m.pgserver.s.ConnectionLimit = m.Config.Postgres.ConnectionLimit + err := m.pgserver.Start(m.Config.Postgres.Addr) + if err != nil { + return errors.Wrap(err, "starting postgres") + } + } + close(m.Started) return nil } @@ -515,6 +539,7 @@ func (m *Command) Close() error { eg.Go(m.Handler.Close) eg.Go(m.Server.Close) eg.Go(m.API.Close) + eg.Go(m.pgserver.Close) if m.gossipMemberSet != nil { eg.Go(m.gossipMemberSet.Close) }