Browse Source

fix: keep PostgreSQL errors in engine transaction and Sync state (#18)

Danilo Fragoso 2 ngày trước cách đây
mục cha
commit
102a03a568

+ 44 - 0
PGPROXY_ERRORS.md

@@ -0,0 +1,44 @@
+# Managed proxy error relay (version 1)
+
+API-key scope validation happens in the managed pgproxy, but PostgreSQL
+transaction and extended-protocol recovery state belongs to the engine. A proxy
+must not generate its own `ReadyForQuery('I')` after rejecting a statement: the
+backend may still have an active transaction or outstanding extended messages.
+
+## Discovery and frame
+
+Supporting engines advertise startup ParameterStatus
+`pizzasql.proxy_errors = 1` and `protocol.proxy_error_relay: true` in
+`GET /capabilities`. The proxy uses this extension only after discovery.
+
+Replace a rejected frontend frame, in the same stream position, with:
+
+- Message type: `y` (private extension, not a PostgreSQL message).
+- Big-endian length: `5` (four length bytes plus one payload byte).
+- Payload: `Q` for a rejected simple Query, or `P` for a rejected Parse.
+
+No SQL, parameters, key scopes or credentials are included in this frame. The
+engine returns SQLSTATE `42501` with a generic permission-denied message. An
+already failed transaction stays failed and returns `25P02`.
+
+For `Q`, the engine sends the error followed by its real ReadyForQuery state.
+For `P`, it fails the extended cycle and discards subsequent frontend messages
+until Sync, where it sends its real ReadyForQuery state. Flush and Terminate
+remain usable during recovery. An error inside a transaction marks it failed;
+ROLLBACK or ROLLBACK TO SAVEPOINT recovers it normally.
+
+The frame cannot execute SQL or grant permissions. The engine listener remains
+trusted loopback-only; this extension does not add authentication to that
+listener. Ordinary clients need not send it.
+
+## Mixed-version rollout
+
+An updated proxy must not send `y` to an engine that did not advertise version
+1. It instead fails closed with FATAL `42501` and closes the tunnel. Closing the
+backend rolls back any uncommitted transaction. The client must reconnect;
+there is no synthetic ReadyForQuery on that fallback path.
+
+Deploying a proxy does not promote tenant engines. Upgrade the approved engine
+pair through the release catalog before expecting recoverable scope errors on
+that tenant. Engine registration/promotion and app deployment remain separate
+operator-approved operations.

+ 3 - 0
pkg/capabilities/capabilities.go

@@ -52,6 +52,7 @@ type Protocol struct {
 	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"`
 }
 
@@ -142,6 +143,7 @@ func Current() *Document {
 			SQLiteCompat:   "3.35.0",
 			SimpleQuery:    true,
 			ExtendedQuery:  true,
+			ProxyErrors:    true,
 			TLS:            "terminated by the managed pgproxy, not the engine listener",
 		},
 
@@ -218,6 +220,7 @@ func Current() *Document {
 			"23505": "unique_violation",
 			"40001": "serialization_failure",
 			"0A000": "feature_not_supported",
+			"42501": "insufficient_privilege (managed proxy)",
 			"42601": "syntax_error",
 			"25P02": "in_failed_sql_transaction",
 			"08P01": "protocol_violation (managed proxy: TLS required)",

+ 39 - 3
pkg/pgserver/connection.go

@@ -171,6 +171,9 @@ func (c *Connection) Handle() error {
 	if err := c.sendParameterStatus("pizzasql.commit", version.Commit); err != nil {
 		return err
 	}
+	if err := c.sendParameterStatus("pizzasql.proxy_errors", "1"); err != nil {
+		return err
+	}
 	if err := c.sendParameterStatus("server_encoding", "UTF8"); err != nil {
 		return err
 	}
@@ -283,7 +286,12 @@ func (c *Connection) initDatabase(dbName string) error {
 
 // handleMessage processes a client message
 func (c *Connection) handleMessage(msg *Message) error {
+	if c.extendedFailed && msg.Type != MsgSync && msg.Type != MsgFlush && msg.Type != MsgTerminate {
+		return nil
+	}
 	switch msg.Type {
+	case MsgProxyError:
+		return c.handleProxyError(msg)
 	case MsgQuery:
 		return c.handleQuery(msg)
 
@@ -390,9 +398,9 @@ func (c *Connection) handleQuery(msg *Message) error {
 		code := ErrCodeSyntaxError
 		if errors.Is(err, pgerror.ErrFeatureNotSupported) {
 			code = ErrCodeFeatureNotSupported
-			if c.txStatus == TxStatusInBlock {
-				c.txStatus = TxStatusFailed
-			}
+		}
+		if c.txStatus == TxStatusInBlock {
+			c.txStatus = TxStatusFailed
 		}
 		c.sendError("ERROR", code, fmt.Sprintf("Syntax error: %v", err))
 		return c.sendReadyForQuery()
@@ -906,9 +914,34 @@ func (c *Connection) handleClose(msg *Message) error {
 
 func (c *Connection) failExtended(code string, err error) error {
 	c.extendedFailed = true
+	if c.txStatus == TxStatusInBlock {
+		c.txStatus = TxStatusFailed
+	}
 	return c.sendError("ERROR", code, err.Error())
 }
 
+func (c *Connection) handleProxyError(msg *Message) error {
+	if len(msg.Data) != 1 || (msg.Data[0] != MsgQuery && msg.Data[0] != MsgParse) {
+		_ = c.sendError("FATAL", ErrCodeProtocolViolation, "invalid proxy error relay")
+		return io.EOF
+	}
+	code, message := "42501", "permission denied by API key scopes"
+	if c.txStatus == TxStatusFailed {
+		code, message = ErrCodeTransactionAborted, "current transaction is aborted, commands ignored until end of transaction block"
+	}
+	if msg.Data[0] == MsgParse {
+		c.extendedCycle = true
+		return c.failExtended(code, errors.New(message))
+	}
+	if c.txStatus == TxStatusInBlock {
+		c.txStatus = TxStatusFailed
+	}
+	if err := c.sendError("ERROR", code, message); err != nil {
+		return err
+	}
+	return c.sendReadyForQuery()
+}
+
 type boundParameter struct {
 	value  []byte
 	oid    int32
@@ -1284,6 +1317,9 @@ func (c *Connection) sendCommandComplete(tag string) error {
 
 // sendError sends an error response
 func (c *Connection) sendError(severity, code, message string) error {
+	if severity == "ERROR" && c.txStatus == TxStatusInBlock {
+		c.txStatus = TxStatusFailed
+	}
 	mb := NewMessageBuilder()
 	mb.AppendByte(ErrorFieldSeverity)
 	mb.WriteString(severity)

+ 14 - 11
pkg/pgserver/protocol.go

@@ -18,17 +18,20 @@ const (
 // Message type constants (first byte of message)
 const (
 	// Frontend (client) messages
-	MsgStartup   = 0   // Startup message (no type byte)
-	MsgQuery     = 'Q' // Simple query
-	MsgTerminate = 'X' // Terminate
-	MsgPassword  = 'p' // Password message
-	MsgParse     = 'P' // Parse (prepared statement)
-	MsgBind      = 'B' // Bind
-	MsgDescribe  = 'D' // Describe
-	MsgExecute   = 'E' // Execute
-	MsgSync      = 'S' // Sync
-	MsgFlush     = 'H' // Flush
-	MsgClose     = 'C' // Close
+	MsgStartup = 0   // Startup message (no type byte)
+	MsgQuery   = 'Q' // Simple query
+	// MsgProxyError is a negotiated managed-proxy extension. Its one-byte
+	// payload is Q or P, replacing a denied Query or Parse in stream order.
+	MsgProxyError = 'y'
+	MsgTerminate  = 'X' // Terminate
+	MsgPassword   = 'p' // Password message
+	MsgParse      = 'P' // Parse (prepared statement)
+	MsgBind       = 'B' // Bind
+	MsgDescribe   = 'D' // Describe
+	MsgExecute    = 'E' // Execute
+	MsgSync       = 'S' // Sync
+	MsgFlush      = 'H' // Flush
+	MsgClose      = 'C' // Close
 
 	// Backend (server) messages
 	MsgAuthenticationOk     = 'R' // Authentication request

+ 89 - 0
pkg/pgserver/proxy_error_test.go

@@ -0,0 +1,89 @@
+package pgserver
+
+import (
+	"bufio"
+	"net"
+	"testing"
+)
+
+func runMessageCycle(t *testing.T, c *Connection, client net.Conn, messages ...*Message) []*Message {
+	t.Helper()
+	done := make(chan error, 1)
+	go func() {
+		for _, msg := range messages {
+			if err := c.handleMessage(msg); err != nil {
+				done <- err
+				return
+			}
+		}
+		done <- nil
+	}()
+	r := bufio.NewReader(client)
+	var response []*Message
+	for {
+		msg, err := ReadMessage(r)
+		if err != nil {
+			t.Fatal(err)
+		}
+		response = append(response, msg)
+		if msg.Type == MsgReadyForQuery {
+			break
+		}
+	}
+	if err := <-done; err != nil {
+		t.Fatal(err)
+	}
+	return response
+}
+
+func TestSyntaxErrorAbortsSimpleAndExtendedTransaction(t *testing.T) {
+	for _, extended := range []bool{false, true} {
+		c, client := newDataConnection(t)
+		runQuery(t, c, client, &Message{Type: MsgQuery, Data: []byte("BEGIN\x00")})
+		var response []*Message
+		if extended {
+			response = runExtendedExecution(t, c, client, "SELEC invalid")
+		} else {
+			response = runQuery(t, c, client, &Message{Type: MsgQuery, Data: []byte("SELEC invalid\x00")})
+		}
+		if errorCode(response[0]) != "42601" || c.txStatus != TxStatusFailed {
+			t.Fatalf("extended=%v did not abort on syntax error", extended)
+		}
+		response = runQuery(t, c, client, &Message{Type: MsgQuery, Data: []byte("SELECT 1\x00")})
+		if errorCode(response[0]) != "25P02" {
+			t.Fatal("failed transaction accepted SELECT")
+		}
+		runQuery(t, c, client, &Message{Type: MsgQuery, Data: []byte("ROLLBACK\x00")})
+	}
+}
+
+func TestProxyErrorRelayTransactionAndSyncRecovery(t *testing.T) {
+	for _, extended := range []bool{false, true} {
+		c, client := newDataConnection(t)
+		for _, sql := range []string{"CREATE TABLE t (id INTEGER PRIMARY KEY)", "BEGIN", "INSERT INTO t VALUES (1)", "SAVEPOINT s"} {
+			runQuery(t, c, client, &Message{Type: MsgQuery, Data: []byte(sql + "\x00")})
+		}
+		kind := byte(MsgQuery)
+		if extended {
+			kind = MsgParse
+		}
+		messages := []*Message{{Type: MsgProxyError, Data: []byte{kind}}}
+		if extended {
+			c.portals[""] = &portal{query: "INSERT INTO t VALUES (2)"}
+			messages = append(messages, &Message{Type: MsgExecute, Data: []byte{0, 0, 0, 0, 0}}, &Message{Type: MsgSync})
+		}
+		response := runMessageCycle(t, c, client, messages...)
+		if errorCode(response[0]) != "42501" {
+			t.Fatalf("extended=%v missing permission error", extended)
+		}
+		if last := response[len(response)-1]; last.Type != MsgReadyForQuery || string(last.Data) != "E" {
+			t.Fatalf("extended=%v false ready state: %v", extended, last)
+		}
+		runQuery(t, c, client, &Message{Type: MsgQuery, Data: []byte("ROLLBACK TO SAVEPOINT s\x00")})
+		response = runQuery(t, c, client, &Message{Type: MsgQuery, Data: []byte("SELECT COUNT(*) FROM t\x00")})
+		if row := dataRowValues(t, response[1]); len(row) != 1 || row[0] != "1" {
+			t.Fatalf("relay lost earlier writes or executed discarded messages: %v", row)
+		}
+		runQuery(t, c, client, &Message{Type: MsgQuery, Data: []byte("ROLLBACK\x00")})
+	}
+}