Query runner: Block known problematic functions if found in parse tree

go.mod

@@ -31,7 +31,7 @@ require (
 	golang.org/x/crypto v0.31.0 // indirect
 	golang.org/x/net v0.33.0
 	google.golang.org/api v0.143.0
-	google.golang.org/protobuf v1.33.0
+	google.golang.org/protobuf v1.36.1
 	gopkg.in/ini.v1 v1.62.0 // indirect
 	gopkg.in/mcuadros/go-syslog.v2 v2.3.0
 )

go.sum

@@ -334,8 +334,8 @@ google.golang.org/protobuf v1.25.0/go.mod h1:9JNX74DMeImyA3h4bdi1ymwjUzf21/xIlba
 google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
 google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
 google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
-google.golang.org/protobuf v1.33.0 h1:uNO2rsAINq/JlFpSdYEKIZ0uKD/R9cpdv0T+yoGwGmI=
-google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
+google.golang.org/protobuf v1.36.1 h1:yBPeRvTftaleIgM3PZ/WBIZ7XM/eEYAaEyCwvyjq/gk=
+google.golang.org/protobuf v1.36.1/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=

input/postgres/explain_analyze.go

@@ -4,11 +4,16 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"slices"
 	"strings"
 
 	"github.com/guregu/null"
 	"github.com/lib/pq"
-	"github.com/pganalyze/collector/util"
+	pg_query "github.com/pganalyze/pg_query_go/v6"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/reflect/protopath"
+	"google.golang.org/protobuf/reflect/protorange"
+	"google.golang.org/protobuf/reflect/protoreflect"
 )
 
 func RunExplainAnalyzeForQueryRun(ctx context.Context, db *sql.DB, query string, parameters []null.String, parameterTypes []string, marker string) (result string, err error) {
@@ -53,16 +58,72 @@ func runExplainAnalyze(ctx context.Context, db *sql.DB, query string, parameters
 }
 
 func validateQuery(query string) error {
-	var isUtil []bool
-	// To be on the safe side never EXPLAIN a statement that can't be parsed,
-	// or multiple statements in one (leading to accidental execution)
-	isUtil, err := util.IsUtilityStmt(query)
-	if err != nil || len(isUtil) != 1 || isUtil[0] {
-		err = fmt.Errorf("query is not permitted to run (multi-statement or utility command?)")
-		return err
+	parseResult, err := pg_query.Parse(query)
+	if err != nil {
+		return fmt.Errorf("query is not permitted to run - failed to parse")
+	}
+	if len(parseResult.Stmts) != 1 {
+		return fmt.Errorf("query is not permitted to run - multi-statement query string")
 	}
 
-	// TODO: Consider adding additional checks here (e.g. blocking known bad function calls)
+	stmt := parseResult.Stmts[0].Stmt.Node
+	switch stmt.(type) {
+	case *pg_query.Node_SelectStmt:
+		// Allowed, continue
+		// Note that we permit wCTEs here (for now), and instead rely on the read-only transaction to block them
+	case *pg_query.Node_InsertStmt, *pg_query.Node_UpdateStmt, *pg_query.Node_DeleteStmt:
+		return fmt.Errorf("query is not permitted to run - DML statement")
+	default:
+		return fmt.Errorf("query is not permitted to run - utility statement")
+	}
+
+	err = validateBlockedFunctions(parseResult)
+	if err != nil {
+		return err
+	}
 
 	return nil
 }
+
+var blockedFunctions = []string{
+	// Blocked because these functions allow exfiltrating data to external servers
+	"dblink",
+	"dblink_connect",
+	"dblink_exec",
+	// Blocked because these functions allow executing arbitrary SQL as input (which can workaround other checks)
+	"xpath_table",
+}
+
+func validateBlockedFunctions(parseResult *pg_query.ParseResult) error {
+	return walkParseTree(parseResult, func(nodeType string, node proto.Message) error {
+		if nodeType != "FuncCall" {
+			return nil
+		}
+
+		f := node.(*pg_query.FuncCall)
+		// The funcname field can be optionally schema qualified, so we take the last item in the list of names
+		nameNode := f.Funcname[len(f.Funcname)-1]
+		name := nameNode.GetString_().Sval
+
+		if slices.Contains(blockedFunctions, name) {
+			return fmt.Errorf("query is not permitted to run - function not allowed: %s", name)
+		}
+		return nil
+	})
+}
+
+type treeWalker func(nodeType string, node proto.Message) error
+
+func walkParseTree(parseResult *pg_query.ParseResult, fn treeWalker) error {
+	return protorange.Range(parseResult.ProtoReflect(), func(p protopath.Values) error {
+		last := p.Index(-1)
+		m, ok := last.Value.Interface().(protoreflect.Message)
+		if ok {
+			err := fn(string(m.Descriptor().Name()), m.Interface())
+			if err != nil {
+				return err
+			}
+		}
+		return nil
+	})
+}

input/postgres/explain_analyze_test.go

@@ -234,14 +234,21 @@ var queryRunTests = []queryRunTestpair{
 		[]null.String{},
 		[]string{},
 		"",
-		"pq: cannot execute UPDATE in a read-only transaction",
+		"query is not permitted to run - DML statement",
 	},
 	{
 		"SELECT 1; UPDATE test SET id = 123",
 		[]null.String{},
 		[]string{},
 		"",
-		"query is not permitted to run (multi-statement or utility command?)",
+		"query is not permitted to run - multi-statement query string",
+	},
+	{
+		"SELECT dblink_exec('host=myhost user=myuser password=mypass dbname=mydb', dblink_build_sql_insert('secret_table', '1', 1, '{\"1\"}', '{\"1\"}'))",
+		[]null.String{},
+		[]string{},
+		"",
+		"query is not permitted to run - function not allowed: dblink_exec",
 	},
 	{
 		"SELECT $1",