fix: escape policy role names and function/procedure names in generated SQL

internal/migration_acceptance_tests/function_cases_test.go

@@ -573,6 +573,50 @@ var functionAcceptanceTestCases = []acceptanceTestCase{
 		},
 		expectedHazardTypes: []diff.MigrationHazardType{diff.MigrationHazardTypeHasUntrackableDependencies},
 	},
+	{
+		name:         "Create function with name containing double quote",
+		oldSchemaDDL: nil,
+		newSchemaDDL: []string{
+			`
+            CREATE FUNCTION "evil""func"(a integer) RETURNS integer
+                LANGUAGE SQL
+                IMMUTABLE
+                RETURNS NULL ON NULL INPUT
+                RETURN a + 1;
+			`,
+		},
+	},
+	{
+		name: "Drop function with name containing double quote",
+		oldSchemaDDL: []string{
+			`
+            CREATE FUNCTION "evil""func"(a integer) RETURNS integer
+                LANGUAGE SQL
+                IMMUTABLE
+                RETURNS NULL ON NULL INPUT
+                RETURN a + 1;
+			`,
+		},
+		newSchemaDDL: nil,
+	},
+	{
+		name: "Create function with SQL injection in name",
+		oldSchemaDDL: []string{
+			`
+            CREATE TABLE foo(id integer);
+			`,
+		},
+		newSchemaDDL: []string{
+			`
+            CREATE TABLE foo(id integer);
+            CREATE FUNCTION "x""; DROP TABLE foo; --"(a integer) RETURNS integer
+                LANGUAGE SQL
+                IMMUTABLE
+                RETURNS NULL ON NULL INPUT
+                RETURN a;
+			`,
+		},
+	},
 }
 
 func TestFunctionTestCases(t *testing.T) {

internal/migration_acceptance_tests/policy_cases_test.go

@@ -658,6 +658,84 @@ var policyAcceptanceTestCases = []acceptanceTestCase{
 
 		expectedPlanErrorIs: diff.ErrNotImplemented,
 	},
+	{
+		name: "Create policy with role containing double quote",
+		roles: []string{
+			`"evil""role"`,
+		},
+		oldSchemaDDL: []string{
+			`
+                CREATE TABLE foobar();
+			`,
+		},
+		newSchemaDDL: []string{
+			`
+                CREATE TABLE foobar();
+                CREATE POLICY foobar_policy ON foobar
+                    AS PERMISSIVE
+                    FOR ALL
+                    TO "evil""role"
+                    USING (true)
+                    WITH CHECK (true);
+			`,
+		},
+		expectedHazardTypes: []diff.MigrationHazardType{
+			diff.MigrationHazardTypeAuthzUpdate,
+		},
+	},
+	{
+		name: "Alter policy with role containing SQL injection payload",
+		roles: []string{
+			`"x""; DROP TABLE foobar; --"`,
+		},
+		oldSchemaDDL: []string{
+			`
+                CREATE TABLE foobar();
+                CREATE POLICY foobar_policy ON foobar
+                    AS PERMISSIVE
+                    FOR ALL
+                    TO PUBLIC
+                    USING (true)
+                    WITH CHECK (true);
+			`,
+		},
+		newSchemaDDL: []string{
+			`
+                CREATE TABLE foobar();
+                CREATE POLICY foobar_policy ON foobar
+                    AS PERMISSIVE
+                    FOR ALL
+                    TO "x""; DROP TABLE foobar; --"
+                    USING (true)
+                    WITH CHECK (true);
+			`,
+		},
+		expectedHazardTypes: []diff.MigrationHazardType{
+			diff.MigrationHazardTypeAuthzUpdate,
+		},
+	},
+	{
+		name: "Create policy with PUBLIC keyword unquoted",
+		oldSchemaDDL: []string{
+			`
+                CREATE TABLE foobar();
+			`,
+		},
+		newSchemaDDL: []string{
+			`
+                CREATE TABLE foobar();
+                CREATE POLICY foobar_policy ON foobar
+                    AS PERMISSIVE
+                    FOR ALL
+                    TO PUBLIC
+                    USING (true)
+                    WITH CHECK (true);
+			`,
+		},
+		expectedHazardTypes: []diff.MigrationHazardType{
+			diff.MigrationHazardTypeAuthzUpdate,
+		},
+	},
 }
 
 func TestPolicyCases(t *testing.T) {

internal/schema/escape_test.go

@@ -58,3 +58,41 @@ func TestEscapeLiteral(t *testing.T) {
 		})
 	}
 }
+
+func TestBuildProcName(t *testing.T) {
+	for _, tc := range []struct {
+		name            string
+		procName        string
+		identityArgs    string
+		schemaName      string
+		expectedEscaped string
+	}{
+		{
+			name:            "simple function",
+			procName:        "add",
+			identityArgs:    "integer, integer",
+			schemaName:      "public",
+			expectedEscaped: `"add"(integer, integer)`,
+		},
+		{
+			name:            "function name with embedded double quote",
+			procName:        `evil"func`,
+			identityArgs:    "integer",
+			schemaName:      "public",
+			expectedEscaped: `"evil""func"(integer)`,
+		},
+		{
+			name:            "injection attempt via double quote",
+			procName:        `x"; DROP TABLE foo; --`,
+			identityArgs:    "",
+			schemaName:      "public",
+			expectedEscaped: `"x""; DROP TABLE foo; --"()`,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			result := buildProcName(tc.procName, tc.identityArgs, tc.schemaName)
+			assert.Equal(t, tc.schemaName, result.SchemaName)
+			assert.Equal(t, tc.expectedEscaped, result.EscapedName)
+		})
+	}
+}

internal/schema/schema.go

@@ -1588,7 +1588,7 @@ func parseJSONTableDependencies(vals []string) ([]TableDependency, error) {
 func buildProcName(name, identityArguments, schemaName string) SchemaQualifiedName {
 	return SchemaQualifiedName{
 		SchemaName:  schemaName,
-		EscapedName: fmt.Sprintf("\"%s\"(%s)", name, identityArguments),
+		EscapedName: fmt.Sprintf("%s(%s)", EscapeIdentifier(name), identityArguments),
 	}
 }
 

pkg/diff/policy_sql_generator.go

@@ -100,6 +100,20 @@ func unforceRLSForTable(t schema.Table) Statement {
 	}
 }
 
+// escapeRoleNames escapes role names for use in SQL policy statements.
+// PUBLIC is a SQL keyword (not an identifier) and must not be quoted.
+func escapeRoleNames(roles []string) []string {
+	escaped := make([]string, len(roles))
+	for i, role := range roles {
+		if role == "PUBLIC" {
+			escaped[i] = role
+		} else {
+			escaped[i] = schema.EscapeIdentifier(role)
+		}
+	}
+	return escaped
+}
+
 type policyDiff struct {
 	oldAndNew[schema.Policy]
 }
@@ -164,7 +178,7 @@ func (psg *policySQLVertexGenerator) Add(p schema.Policy) ([]Statement, error) {
 	}
 	sb.WriteString(fmt.Sprintf("\n\tFOR %s", cmdSQL))
 
-	sb.WriteString(fmt.Sprintf("\n\tTO %s", strings.Join(p.AppliesTo, ", ")))
+	sb.WriteString(fmt.Sprintf("\n\tTO %s", strings.Join(escapeRoleNames(p.AppliesTo), ", ")))
 
 	if p.UsingExpression != "" {
 		sb.WriteString(fmt.Sprintf("\n\tUSING (%s)", p.UsingExpression))
@@ -223,7 +237,7 @@ func (psg *policySQLVertexGenerator) Alter(diff policyDiff) ([]Statement, error)
 	var alterPolicyParts []string
 
 	if !cmp.Equal(oldCopy.AppliesTo, diff.new.AppliesTo) {
-		alterPolicyParts = append(alterPolicyParts, fmt.Sprintf("TO %s", strings.Join(diff.new.AppliesTo, ", ")))
+		alterPolicyParts = append(alterPolicyParts, fmt.Sprintf("TO %s", strings.Join(escapeRoleNames(diff.new.AppliesTo), ", ")))
 		oldCopy.AppliesTo = diff.new.AppliesTo
 	}
 

pkg/diff/policy_sql_generator_test.go

@@ -0,0 +1,94 @@
+package diff
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/stripe/pg-schema-diff/internal/schema"
+)
+
+func TestEscapeRoleNames(t *testing.T) {
+	for _, tc := range []struct {
+		name     string
+		input    []string
+		expected []string
+	}{
+		{
+			name:     "PUBLIC is not escaped",
+			input:    []string{"PUBLIC"},
+			expected: []string{"PUBLIC"},
+		},
+		{
+			name:     "simple role is escaped",
+			input:    []string{"my_role"},
+			expected: []string{`"my_role"`},
+		},
+		{
+			name:     "mixed PUBLIC and regular roles",
+			input:    []string{"PUBLIC", "role_1", "role_2"},
+			expected: []string{"PUBLIC", `"role_1"`, `"role_2"`},
+		},
+		{
+			name:     "role with embedded double quote",
+			input:    []string{`evil"role`},
+			expected: []string{`"evil""role"`},
+		},
+		{
+			name:     "injection attempt via role name",
+			input:    []string{`x"; DROP TABLE foo; --`},
+			expected: []string{`"x""; DROP TABLE foo; --"`},
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			assert.Equal(t, tc.expected, escapeRoleNames(tc.input))
+		})
+	}
+}
+
+func TestPolicySQLGenerator_Add_EscapesRoleNames(t *testing.T) {
+	table := schema.Table{
+		SchemaQualifiedName: schema.SchemaQualifiedName{
+			SchemaName:  "public",
+			EscapedName: `"foo"`,
+		},
+	}
+	psg, err := newPolicySQLVertexGenerator(nil, table)
+	require.NoError(t, err)
+
+	for _, tc := range []struct {
+		name        string
+		appliesTo   []string
+		expectedSub string
+	}{
+		{
+			name:        "PUBLIC unquoted",
+			appliesTo:   []string{"PUBLIC"},
+			expectedSub: "\n\tTO PUBLIC",
+		},
+		{
+			name:        "regular role quoted",
+			appliesTo:   []string{"my_role"},
+			expectedSub: "\n\tTO \"my_role\"",
+		},
+		{
+			name:        "role with embedded double quote",
+			appliesTo:   []string{`evil"role`},
+			expectedSub: "\n\tTO \"evil\"\"role\"",
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			graph, err := psg.Add(schema.Policy{
+				EscapedName:  `"test_policy"`,
+				IsPermissive: true,
+				AppliesTo:    tc.appliesTo,
+				Cmd:          schema.AllPolicyCmd,
+			})
+			require.NoError(t, err)
+			require.NotEmpty(t, graph.vertices)
+			require.NotEmpty(t, graph.vertices[0].statements)
+			ddl := graph.vertices[0].statements[0].DDL
+			assert.Contains(t, ddl, tc.expectedSub)
+		})
+	}
+}