Преглед на файлове

fix: bind parameters outside comments and exactly match system probes (#23)

Danilo Fragoso преди 2 дни
родител
ревизия
39fcea74de
променени са 4 файла, в които са добавени 278 реда и са изтрити 76 реда
  1. 1 1
      README.md
  2. 53 75
      pkg/pgserver/connection.go
  3. 116 0
      pkg/pgserver/sql_matching.go
  4. 108 0
      pkg/pgserver/sql_matching_test.go

+ 1 - 1
README.md

@@ -320,7 +320,7 @@ pgcli postgresql://localhost:5432/pizzasql
 1. **SQL Dialect**: PizzaSQL uses **SQLite syntax**, not PostgreSQL syntax
    - 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)
+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.
 
 ---

+ 53 - 75
pkg/pgserver/connection.go

@@ -349,7 +349,6 @@ func (c *Connection) handleQuery(msg *Message) error {
 	}
 
 	// Handle special PostgreSQL system queries that drivers send
-	sqlUpper := strings.ToUpper(strings.TrimSpace(sql))
 	singleStatement := isSingleStatementSQL(sql)
 
 	// In a failed transaction only ROLLBACK (or ROLLBACK TO SAVEPOINT) is
@@ -361,22 +360,16 @@ func (c *Connection) handleQuery(msg *Message) error {
 		return c.sendReadyForQuery()
 	}
 
-	// lib/pq and other drivers query these for connection validation
-	if singleStatement && strings.Contains(sqlUpper, "SELECT VERSION()") {
-		// Return a fake PostgreSQL version
-		return c.handleVersionQuery()
-	}
-
-	if singleStatement && strings.Contains(sqlUpper, "SELECT CURRENT_USER") {
-		// Return the current user
-		return c.handleCurrentUserQuery()
-	}
-
-	if singleStatement && strings.Contains(sqlUpper, "SHOW") && (strings.Contains(sqlUpper, "SERVER_VERSION") ||
-		strings.Contains(sqlUpper, "SERVER_ENCODING") ||
-		strings.Contains(sqlUpper, "CLIENT_ENCODING")) {
-		// Handle SHOW commands
-		return c.handleShowCommand(sqlUpper)
+	// Recognize only complete initialization-query token sequences.
+	if singleStatement {
+		switch kind := systemQueryKind(sql); kind {
+		case "version":
+			return c.handleVersionQuery()
+		case "current_user":
+			return c.handleCurrentUserQuery()
+		case "server_version", "server_encoding", "client_encoding":
+			return c.handleShowCommand(strings.ToUpper(kind))
+		}
 	}
 	if result, handled, err := c.catalogResult(sql); handled {
 		if err != nil {
@@ -670,18 +663,7 @@ func (c *Connection) catalogResult(sql string) (*executor.Result, bool, error) {
 	if !isSingleStatementSQL(sql) {
 		return nil, false, nil
 	}
-	upper := strings.ToUpper(sql)
-	var source string
-	for _, candidate := range []string{
-		"INFORMATION_SCHEMA.TABLES", "INFORMATION_SCHEMA.COLUMNS",
-		"INFORMATION_SCHEMA.TABLE_CONSTRAINTS", "INFORMATION_SCHEMA.KEY_COLUMN_USAGE",
-		"PG_TABLES", "PG_INDEXES",
-	} {
-		if strings.Contains(upper, "FROM "+candidate) {
-			source = candidate
-			break
-		}
-	}
+	source := catalogSource(sql)
 	if source == "" {
 		return nil, false, nil
 	}
@@ -801,35 +783,6 @@ func (c *Connection) catalogResult(sql string) (*executor.Result, bool, error) {
 	return result, true, nil
 }
 
-func isSingleStatementSQL(sql string) bool {
-	trimmed := strings.TrimSpace(sql)
-	if strings.HasSuffix(trimmed, ";") {
-		trimmed = strings.TrimSpace(strings.TrimSuffix(trimmed, ";"))
-	}
-	inString, inIdent := false, false
-	for i := 0; i < len(trimmed); i++ {
-		switch trimmed[i] {
-		case '\'':
-			if !inIdent {
-				if inString && i+1 < len(trimmed) && trimmed[i+1] == '\'' {
-					i++
-					continue
-				}
-				inString = !inString
-			}
-		case '"':
-			if !inString {
-				inIdent = !inIdent
-			}
-		case ';':
-			if !inString && !inIdent {
-				return false
-			}
-		}
-	}
-	return true
-}
-
 // isRollbackStatement reports whether sql parses as a single ROLLBACK
 // statement, including ROLLBACK TO SAVEPOINT.
 func isRollbackStatement(sql string) bool {
@@ -996,34 +949,59 @@ func parameterFormat(formats []int16, index int) int16 {
 
 func bindQuery(query string, params []boundParameter) (string, error) {
 	var result strings.Builder
-	inString, inIdent := false, false
+	var quote byte
 	for i := 0; i < len(query); {
 		ch := query[i]
-		if ch == '\'' && !inIdent {
+		if quote != 0 {
 			result.WriteByte(ch)
-			if inString && i+1 < len(query) && query[i+1] == '\'' {
-				result.WriteByte(query[i+1])
-				i += 2
-				continue
-			}
-			inString = !inString
 			i++
+			if ch == quote {
+				if i < len(query) && query[i] == quote {
+					result.WriteByte(query[i])
+					i++
+				} else {
+					quote = 0
+				}
+			}
 			continue
 		}
-		if ch == '"' && !inString {
-			inIdent = !inIdent
+		if i+1 < len(query) && ch == '-' && query[i+1] == '-' {
+			end := strings.IndexByte(query[i:], '\n')
+			if end < 0 {
+				result.WriteString(query[i:])
+				break
+			}
+			result.WriteString(query[i : i+end+1])
+			i += end + 1
+			continue
+		}
+		if i+1 < len(query) && ch == '/' && query[i+1] == '*' {
+			end := strings.Index(query[i+2:], "*/")
+			if end < 0 {
+				result.WriteString(query[i:])
+				break
+			}
+			result.WriteString(query[i : i+end+4])
+			i += end + 4
+			continue
+		}
+		if ch == '\'' || ch == '"' || ch == '`' || ch == '[' {
+			quote = ch
+			if ch == '[' {
+				quote = ']'
+			}
 			result.WriteByte(ch)
 			i++
 			continue
 		}
-		if ch == '$' && !inString && !inIdent && i+1 < len(query) && query[i+1] >= '0' && query[i+1] <= '9' {
+		if ch == '$' && i+1 < len(query) && query[i+1] >= '0' && query[i+1] <= '9' {
 			end := i + 1
 			for end < len(query) && query[end] >= '0' && query[end] <= '9' {
 				end++
 			}
-			n, _ := strconv.Atoi(query[i+1 : end])
-			if n < 1 || n > len(params) {
-				return "", fmt.Errorf("parameter $%d was not provided", n)
+			n, err := strconv.Atoi(query[i+1 : end])
+			if err != nil || n < 1 || n > len(params) {
+				return "", fmt.Errorf("parameter %s was not provided", query[i:end])
 			}
 			literal, err := parameterLiteral(params[n-1])
 			if err != nil {
@@ -1467,13 +1445,13 @@ func (c *Connection) handleShowCommand(sqlUpper string) error {
 	var value string
 	var name string
 
-	if strings.Contains(sqlUpper, "SERVER_VERSION") {
+	if sqlUpper == "SERVER_VERSION" {
 		name = "server_version"
 		value = version.PostgreSQLCompatVersion
-	} else if strings.Contains(sqlUpper, "SERVER_ENCODING") {
+	} else if sqlUpper == "SERVER_ENCODING" {
 		name = "server_encoding"
 		value = "UTF8"
-	} else if strings.Contains(sqlUpper, "CLIENT_ENCODING") {
+	} else if sqlUpper == "CLIENT_ENCODING" {
 		name = "client_encoding"
 		value = "UTF8"
 	} else {

+ 116 - 0
pkg/pgserver/sql_matching.go

@@ -0,0 +1,116 @@
+package pgserver
+
+import (
+	"strings"
+
+	"github.com/danfragoso/pizzasql-next/pkg/lexer"
+)
+
+func nextSQLToken(l *lexer.Lexer) lexer.Token {
+	for {
+		token := l.NextToken()
+		if token.Type != lexer.TokenComment {
+			return token
+		}
+	}
+}
+
+// Match complete token sequences, never substrings inside literals/comments or
+// projections that happen to mention a driver initialization query.
+func systemQueryKind(sql string) string {
+	l := lexer.New(sql)
+	first := nextSQLToken(l)
+	if first.Type != lexer.TokenSELECT && !(first.Type == lexer.TokenIdent && strings.EqualFold(first.Literal, "SHOW")) {
+		return ""
+	}
+	tokens := []string{strings.ToUpper(first.Literal)}
+	for len(tokens) <= 4 {
+		token := nextSQLToken(l)
+		if token.Type == lexer.TokenSemicolon {
+			if nextSQLToken(l).Type != lexer.TokenEOF {
+				return ""
+			}
+			break
+		}
+		if token.Type == lexer.TokenEOF {
+			break
+		}
+		if token.Type == lexer.TokenString || token.Type == lexer.TokenBlob || token.Type == lexer.TokenError {
+			return ""
+		}
+		tokens = append(tokens, strings.ToUpper(token.Literal))
+	}
+	switch strings.Join(tokens, " ") {
+	case "SELECT VERSION ( )":
+		return "version"
+	case "SELECT CURRENT_USER":
+		return "current_user"
+	case "SHOW SERVER_VERSION":
+		return "server_version"
+	case "SHOW SERVER_ENCODING":
+		return "server_encoding"
+	case "SHOW CLIENT_ENCODING":
+		return "client_encoding"
+	default:
+		return ""
+	}
+}
+
+func isSingleStatementSQL(sql string) bool {
+	l := lexer.New(sql)
+	ended := false
+	for {
+		token := nextSQLToken(l)
+		if token.Type == lexer.TokenEOF {
+			return true
+		}
+		if token.Type == lexer.TokenSemicolon {
+			ended = true
+			continue
+		}
+		if ended {
+			return false
+		}
+	}
+}
+
+func catalogSource(sql string) string {
+	l := lexer.New(sql)
+	if nextSQLToken(l).Type != lexer.TokenSELECT {
+		return ""
+	}
+	depth := 0
+	for {
+		token := nextSQLToken(l)
+		switch token.Type {
+		case lexer.TokenEOF, lexer.TokenError:
+			return ""
+		case lexer.TokenLParen:
+			depth++
+		case lexer.TokenRParen:
+			depth--
+		case lexer.TokenFROM:
+			if depth != 0 {
+				continue
+			}
+			name := nextSQLToken(l)
+			if name.Type != lexer.TokenIdent {
+				return ""
+			}
+			source := strings.ToUpper(name.Literal)
+			if nextSQLToken(l).Type == lexer.TokenDot {
+				name = nextSQLToken(l)
+				if name.Type != lexer.TokenIdent {
+					return ""
+				}
+				source += "." + strings.ToUpper(name.Literal)
+			}
+			switch source {
+			case "INFORMATION_SCHEMA.TABLES", "INFORMATION_SCHEMA.COLUMNS", "INFORMATION_SCHEMA.TABLE_CONSTRAINTS", "INFORMATION_SCHEMA.KEY_COLUMN_USAGE", "PG_TABLES", "PG_INDEXES":
+				return source
+			default:
+				return ""
+			}
+		}
+	}
+}

+ 108 - 0
pkg/pgserver/sql_matching_test.go

@@ -0,0 +1,108 @@
+package pgserver
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestBindQuerySkipsCommentsAndQuotedRegions(t *testing.T) {
+	for _, tc := range []struct{ query, want string }{
+		{"SELECT $1 /* ' $99 */ + $2", "SELECT 1 /* ' $99 */ + 2"},
+		{"SELECT $1 -- '$99\n + $2", "SELECT 1 -- '$99\n + 2"},
+		{"SELECT $1 -- $99", "SELECT 1 -- $99"},
+		{"SELECT '$99''$88', \"$77\", `$66`, [$55], $1", "SELECT '$99''$88', \"$77\", `$66`, [$55], 1"},
+		{"SELECT '$1 /* no comment */', $2", "SELECT '$1 /* no comment */', 2"},
+	} {
+		got, err := bindQuery(tc.query, []boundParameter{{value: []byte("1")}, {value: []byte("2")}})
+		if err != nil || got != tc.want {
+			t.Fatalf("%q: got %q err=%v want %q", tc.query, got, err, tc.want)
+		}
+	}
+	if _, err := bindQuery("SELECT $99999999999999999999999999", nil); err == nil {
+		t.Fatal("overflowed parameter number accepted")
+	}
+}
+
+func TestCommentedPreparedQueryExecutesBothParameters(t *testing.T) {
+	c, client := newDataConnection(t)
+	c.statements[""] = &preparedStatement{query: "SELECT $1 /* ' $99 */ + $2", paramOIDs: []int32{20, 20}}
+	bind := NewMessageBuilder()
+	bind.WriteString("")
+	bind.WriteString("")
+	bind.WriteInt16(0)
+	bind.WriteInt16(2)
+	for _, value := range []string{"1", "2"} {
+		bind.WriteInt32(int32(len(value)))
+		bind.WriteBytes([]byte(value))
+	}
+	bind.WriteInt16(0)
+	if err := c.handleBind(&Message{Type: MsgBind, Data: bind.Bytes()}); err != nil {
+		t.Fatal(err)
+	}
+	portal := c.portals[""]
+	if portal == nil {
+		t.Fatal("Bind failed")
+	}
+	messages := runExtendedExecution(t, c, client, portal.query)
+	found := false
+	for _, message := range messages {
+		if message.Type == MsgErrorResponse {
+			t.Fatalf("bound execution failed: %s", errorCode(message))
+		}
+		if message.Type == MsgDataRow {
+			values := dataRowValues(t, message)
+			if len(values) != 1 || values[0] != "3" {
+				t.Fatal(values)
+			}
+			found = true
+		}
+	}
+	if !found {
+		t.Fatal("no result row")
+	}
+}
+
+func TestSystemQueryMatchingDoesNotInterceptLiteralsOrProjections(t *testing.T) {
+	c, client := newDataConnection(t)
+	for _, literal := range []string{"SELECT VERSION()", "SELECT CURRENT_USER", "SHOW SERVER_VERSION", "FROM PG_TABLES", "FROM INFORMATION_SCHEMA.TABLES"} {
+		messages := runQuery(t, c, client, &Message{Type: MsgQuery, Data: []byte("SELECT '" + literal + "'\x00")})
+		found := false
+		for _, message := range messages {
+			if message.Type == MsgErrorResponse {
+				t.Fatalf("literal %q failed: %s", literal, errorCode(message))
+			}
+			if message.Type == MsgDataRow {
+				values := dataRowValues(t, message)
+				if len(values) != 1 || values[0] != literal {
+					t.Fatalf("literal %q intercepted: %v", literal, values)
+				}
+				found = true
+			}
+		}
+		if !found {
+			t.Fatalf("literal %q produced no row", literal)
+		}
+	}
+	for _, sql := range []string{"SELECT version() + 1", "SELECT current_user, 1", "SELECT 1 /* SELECT VERSION() */", "SHOW other_server_version"} {
+		if kind := systemQueryKind(sql); kind != "" {
+			t.Fatalf("%q intercepted as %s", sql, kind)
+		}
+	}
+	for _, sql := range []string{"select version ( ) /* ; SELECT 2 */;", "-- initialization\nSELECT version()"} {
+		messages := runQuery(t, c, client, &Message{Type: MsgQuery, Data: []byte(sql + "\x00")})
+		if values := dataRowValues(t, messages[1]); !strings.Contains(values[0], "PizzaSQL") {
+			t.Fatal(values)
+		}
+	}
+}
+
+func TestSingleStatementMatchingIgnoresCommentSemicolons(t *testing.T) {
+	for _, sql := range []string{"SELECT ';'", "SELECT 1 /* ; ' */; -- ; SELECT 2", "SELECT 1; ;"} {
+		if !isSingleStatementSQL(sql) {
+			t.Fatalf("single statement rejected: %q", sql)
+		}
+	}
+	if isSingleStatementSQL("SELECT 1; /* ; */ SELECT 2") {
+		t.Fatal("batch accepted as one statement")
+	}
+}