Răsfoiți Sursa

fix: reject unsupported extended-protocol requests before execution (#25)

Danilo Fragoso 1 zi în urmă
părinte
comite
9a74d2bbff
4 a modificat fișierele cu 175 adăugiri și 24 ștergeri
  1. 2 1
      README.md
  2. 12 6
      pkg/capabilities/capabilities.go
  3. 42 17
      pkg/pgserver/connection.go
  4. 119 0
      pkg/pgserver/protocol_guards_test.go

+ 2 - 1
README.md

@@ -321,7 +321,8 @@ pgcli postgresql://localhost:5432/pizzasql
    - Use `INTEGER PRIMARY KEY AUTOINCREMENT`, not `SERIAL`
    - Use `TEXT`, not `VARCHAR` with enforced length
 2. **Parameter placeholders**: use `$1, $2...` (PostgreSQL style) or `?` (SQLite style). PG Bind rewrites `$n` only outside quoted regions and line/block comments. Initialization probes such as `SELECT version()` and `SHOW server_version` match complete token sequences, never substrings inside SQL literals or unrelated projections.
-3. **No PostgreSQL-specific features**: no schemas, roles, `ARRAY`, `JSONB`, etc.
+3. **Extended protocol scope**: Parse accepts one command; Bind requests must use text parameters/all-text results. Describe, binary results, and nonzero Execute row limits fail with SQLSTATE `0A000` before executing side effects. Portals are one-shot, not resumable. Unsupported binary parameter OIDs also report `0A000`.
+4. **No PostgreSQL-specific features**: no schemas, roles, `ARRAY`, `JSONB`, etc.
 
 ---
 

+ 12 - 6
pkg/capabilities/capabilities.go

@@ -48,12 +48,15 @@ type Document struct {
 }
 
 type Protocol struct {
-	PostgresCompat string `json:"postgres_compat"`
-	SQLiteCompat   string `json:"sqlite_compat"`
-	SimpleQuery    bool   `json:"simple_query"`
-	ExtendedQuery  bool   `json:"extended_query"`
-	ProxyErrors    bool   `json:"proxy_error_relay"`
-	TLS            string `json:"tls"`
+	PostgresCompat string  `json:"postgres_compat"`
+	SQLiteCompat   string  `json:"sqlite_compat"`
+	SimpleQuery    bool    `json:"simple_query"`
+	ExtendedQuery  bool    `json:"extended_query"`
+	Describe       Feature `json:"describe"`
+	BinaryResults  Feature `json:"binary_results"`
+	BoundedExecute Feature `json:"bounded_execute"`
+	ProxyErrors    bool    `json:"proxy_error_relay"`
+	TLS            string  `json:"tls"`
 }
 
 type Notifications struct {
@@ -143,6 +146,9 @@ func Current() *Document {
 			SQLiteCompat:   "3.35.0",
 			SimpleQuery:    true,
 			ExtendedQuery:  true,
+			Describe:       Feature{Supported: false, Status: "rejected", SQLState: "0A000", Notes: "use simple query or Parse/Bind/Execute without Describe; no speculative execution for metadata"},
+			BinaryResults:  Feature{Supported: false, Status: "rejected", SQLState: "0A000", Notes: "all-text result formats only; per-column format lists are also rejected"},
+			BoundedExecute: Feature{Supported: false, Status: "rejected", SQLState: "0A000", Notes: "Execute maxRows must be zero; portals are one-shot, not resumable"},
 			ProxyErrors:    true,
 			TLS:            "terminated by the managed pgproxy, not the engine listener",
 		},

+ 42 - 17
pkg/pgserver/connection.go

@@ -482,6 +482,9 @@ func (c *Connection) handleParse(msg *Message) error {
 		oids[i] = int32(binary.BigEndian.Uint32(msg.Data[pos : pos+4]))
 		pos += 4
 	}
+	if !isSingleStatementSQL(query) {
+		return c.failExtended(ErrCodeSyntaxError, fmt.Errorf("prepared statements must contain one command"))
+	}
 	c.statements[name] = &preparedStatement{query: query, paramOIDs: oids}
 	return c.writeMessage(MsgParseComplete, nil)
 }
@@ -508,6 +511,14 @@ func (c *Connection) handleBind(msg *Message) error {
 	}
 	paramCount := int(binary.BigEndian.Uint16(msg.Data[pos : pos+2]))
 	pos += 2
+	if len(formats) > 1 && len(formats) != paramCount {
+		return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("Bind parameter format count does not match parameters"))
+	}
+	for _, format := range formats {
+		if format != 0 && format != 1 {
+			return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Bind parameter format code"))
+		}
+	}
 	params := make([]boundParameter, paramCount)
 	for i := range params {
 		if pos+4 > len(msg.Data) {
@@ -527,13 +538,28 @@ func (c *Connection) handleBind(msg *Message) error {
 		params[i].value = append([]byte(nil), msg.Data[pos:pos+int(length)]...)
 		pos += int(length)
 	}
-	_, pos, err = readInt16List(msg.Data, pos) // result formats; text output is currently used
+	resultFormats, pos, err := readInt16List(msg.Data, pos)
 	if err != nil || pos != len(msg.Data) {
 		return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Bind result format list"))
 	}
+	for _, format := range resultFormats {
+		if format != 0 && format != 1 {
+			return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Bind result format code"))
+		}
+		if format == 1 {
+			return c.failExtended(ErrCodeFeatureNotSupported, fmt.Errorf("binary result format is not supported; request text results"))
+		}
+	}
+	if len(resultFormats) > 1 {
+		return c.failExtended(ErrCodeFeatureNotSupported, fmt.Errorf("per-column result format lists are not supported; request all-text results"))
+	}
 	query, err := bindQuery(statement.query, params)
 	if err != nil {
-		return c.failExtended(ErrCodeInvalidParameter, err)
+		code := ErrCodeInvalidParameter
+		if errors.Is(err, pgerror.ErrFeatureNotSupported) {
+			code = ErrCodeFeatureNotSupported
+		}
+		return c.failExtended(code, err)
 	}
 	c.portals[portalName] = &portal{query: query}
 	return c.writeMessage(MsgBindComplete, nil)
@@ -552,18 +578,9 @@ func (c *Connection) handleDescribe(msg *Message) error {
 	}
 	switch msg.Data[0] {
 	case 'S':
-		statement, ok := c.statements[name]
-		if !ok {
+		if _, ok := c.statements[name]; !ok {
 			return c.failExtended(ErrCodeInvalidParameter, fmt.Errorf("prepared statement %q does not exist", name))
 		}
-		mb := NewMessageBuilder()
-		mb.WriteInt16(int16(len(statement.paramOIDs)))
-		for _, oid := range statement.paramOIDs {
-			mb.WriteInt32(oid)
-		}
-		if err := c.writeMessage(MsgParameterDescription, mb.Bytes()); err != nil {
-			return err
-		}
 	case 'P':
 		if _, ok := c.portals[name]; !ok {
 			return c.failExtended(ErrCodeInvalidParameter, fmt.Errorf("portal %q does not exist", name))
@@ -571,9 +588,10 @@ func (c *Connection) handleDescribe(msg *Message) error {
 	default:
 		return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Describe target"))
 	}
-	// Execute sends the row description once the bound statement has been
-	// analyzed, avoiding side effects during Describe.
-	return c.writeMessage(MsgNoData, nil)
+	// Do not claim NoData for row-producing statements or execute queries to
+	// discover metadata. Clients must avoid Describe until side-effect-free
+	// parameter/row description is implemented.
+	return c.failExtended(ErrCodeFeatureNotSupported, fmt.Errorf("Describe metadata is not supported; use simple query or Parse/Bind/Execute without Describe"))
 }
 
 func (c *Connection) handleExecute(msg *Message) error {
@@ -591,13 +609,20 @@ func (c *Connection) handleExecute(msg *Message) error {
 	if portal.executed {
 		return c.failExtended(ErrCodeFeatureNotSupported, fmt.Errorf("portal can only be executed once"))
 	}
-	portal.executed = true
+	maxRows := int32(binary.BigEndian.Uint32(msg.Data[pos : pos+4]))
+	if maxRows < 0 {
+		return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Execute row limit"))
+	}
 	// Check the transaction state before catalog emulation. Catalog queries do
 	// not all parse as regular PizzaSQL statements, but must still return 25P02
 	// while the transaction is aborted.
 	if c.txStatus == TxStatusFailed && !isRollbackStatement(portal.query) {
 		return c.failExtended(ErrCodeTransactionAborted, fmt.Errorf("current transaction is aborted, commands ignored until end of transaction block"))
 	}
+	if maxRows != 0 {
+		return c.failExtended(ErrCodeFeatureNotSupported, fmt.Errorf("bounded Execute is not supported; use maxRows=0"))
+	}
+	portal.executed = true
 	if result, handled, catalogErr := c.catalogResult(portal.query); handled {
 		if catalogErr != nil {
 			return c.failExtended(ErrCodeInternalError, catalogErr)
@@ -1051,7 +1076,7 @@ func parameterLiteral(param boundParameter) (string, error) {
 			}
 			return strconv.FormatInt(int64(binary.BigEndian.Uint64(param.value)), 10), nil
 		default:
-			return "", fmt.Errorf("binary format is unsupported for OID %d", param.oid)
+			return "", fmt.Errorf("%w: binary parameter format for OID %d", pgerror.ErrFeatureNotSupported, param.oid)
 		}
 	}
 	value := string(param.value)

+ 119 - 0
pkg/pgserver/protocol_guards_test.go

@@ -0,0 +1,119 @@
+package pgserver
+
+import (
+	"bufio"
+	"testing"
+)
+
+func guardConnection(t *testing.T) (*Connection, *countingConn) {
+	t.Helper()
+	c, _ := newDataConnection(t)
+	conn := &countingConn{}
+	c.conn = conn
+	c.writer = bufio.NewWriter(conn)
+	if _, err := c.executor.Execute(parseStmt(t, "CREATE TABLE t (id INTEGER PRIMARY KEY)")); err != nil {
+		t.Fatal(err)
+	}
+	return c, conn
+}
+
+func guardBind(parameterFormats, resultFormats []int16) []byte {
+	b := NewMessageBuilder()
+	b.WriteString("p")
+	b.WriteString("q")
+	b.WriteInt16(int16(len(parameterFormats)))
+	for _, format := range parameterFormats {
+		b.WriteInt16(format)
+	}
+	b.WriteInt16(0)
+	b.WriteInt16(int16(len(resultFormats)))
+	for _, format := range resultFormats {
+		b.WriteInt16(format)
+	}
+	return b.Bytes()
+}
+
+func TestUnsupportedExtendedRequestsFailBeforeSideEffects(t *testing.T) {
+	for _, tc := range []struct {
+		name    string
+		message *Message
+		code    string
+	}{
+		{"binary_results", &Message{Type: MsgBind, Data: guardBind(nil, []int16{1})}, ErrCodeFeatureNotSupported},
+		{"per_column_results", &Message{Type: MsgBind, Data: guardBind(nil, []int16{0, 0})}, ErrCodeFeatureNotSupported},
+		{"invalid_result_format", &Message{Type: MsgBind, Data: guardBind(nil, []int16{2})}, ErrCodeProtocolViolation},
+		{"invalid_parameter_format", &Message{Type: MsgBind, Data: guardBind([]int16{2}, nil)}, ErrCodeProtocolViolation},
+		{"parameter_format_count", &Message{Type: MsgBind, Data: guardBind([]int16{0, 0}, nil)}, ErrCodeProtocolViolation},
+		{"describe_statement", &Message{Type: MsgDescribe, Data: []byte{'S', 'q', 0}}, ErrCodeFeatureNotSupported},
+		{"describe_portal", &Message{Type: MsgDescribe, Data: []byte{'P', 'p', 0}}, ErrCodeFeatureNotSupported},
+		{"bounded_execute", &Message{Type: MsgExecute, Data: []byte{'p', 0, 0, 0, 0, 1}}, ErrCodeFeatureNotSupported},
+		{"negative_execute", &Message{Type: MsgExecute, Data: []byte{'p', 0, 255, 255, 255, 255}}, ErrCodeProtocolViolation},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			c, conn := guardConnection(t)
+			query := "INSERT INTO t VALUES (1) RETURNING id"
+			c.statements["q"] = &preparedStatement{query: query}
+			c.portals["p"] = &portal{query: query}
+			if err := c.handleMessage(tc.message); err != nil {
+				t.Fatal(err)
+			}
+			// Requests pipelined after the error must be discarded until Sync.
+			if err := c.handleMessage(&Message{Type: MsgExecute, Data: []byte{'p', 0, 0, 0, 0, 0}}); err != nil {
+				t.Fatal(err)
+			}
+			if err := c.handleMessage(&Message{Type: MsgSync}); err != nil {
+				t.Fatal(err)
+			}
+			messages := readAllMessages(t, conn.bytes())
+			if len(messages) != 2 || errorCode(messages[0]) != tc.code || messages[1].Type != MsgReadyForQuery {
+				t.Fatalf("unexpected recovery: %v", messages)
+			}
+			if result, err := c.executor.Execute(parseStmt(t, "SELECT COUNT(*) FROM t")); err != nil || result.Rows[0][0] != int64(0) {
+				t.Fatalf("rejected request had side effects: %v %v", result, err)
+			}
+		})
+	}
+}
+
+func TestDescribeFailurePreservesFailedTransactionState(t *testing.T) {
+	c, conn := guardConnection(t)
+	if err := c.handleMessage(&Message{Type: MsgQuery, Data: []byte("BEGIN\x00")}); err != nil {
+		t.Fatal(err)
+	}
+	c.portals["p"] = &portal{query: "INSERT INTO t VALUES (1) RETURNING id"}
+	if err := c.handleMessage(&Message{Type: MsgDescribe, Data: []byte{'P', 'p', 0}}); err != nil {
+		t.Fatal(err)
+	}
+	if err := c.handleMessage(&Message{Type: MsgSync}); err != nil {
+		t.Fatal(err)
+	}
+	messages := readAllMessages(t, conn.bytes())
+	last := messages[len(messages)-1]
+	if last.Type != MsgReadyForQuery || len(last.Data) != 1 || last.Data[0] != TxStatusFailed {
+		t.Fatalf("Describe did not fail transaction: %v", last)
+	}
+	if err := c.handleMessage(&Message{Type: MsgQuery, Data: []byte("ROLLBACK\x00")}); err != nil {
+		t.Fatal(err)
+	}
+	if c.txStatus != TxStatusIdle {
+		t.Fatal("rollback did not recover")
+	}
+}
+
+func TestParseRejectsMultipleCommands(t *testing.T) {
+	c, conn := guardConnection(t)
+	parse := NewMessageBuilder()
+	parse.WriteString("q")
+	parse.WriteString("SELECT 1; INSERT INTO t VALUES (1)")
+	parse.WriteInt16(0)
+	if err := c.handleMessage(&Message{Type: MsgParse, Data: parse.Bytes()}); err != nil {
+		t.Fatal(err)
+	}
+	if _, ok := c.statements["q"]; ok {
+		t.Fatal("multiple commands stored")
+	}
+	messages := readAllMessages(t, conn.bytes())
+	if len(messages) != 1 || errorCode(messages[0]) != ErrCodeSyntaxError {
+		t.Fatal(messages)
+	}
+}