diff --git a/internal/migration_acceptance_tests/enum_cases_test.go b/internal/migration_acceptance_tests/enum_cases_test.go index cd161e8..667f359 100644 --- a/internal/migration_acceptance_tests/enum_cases_test.go +++ b/internal/migration_acceptance_tests/enum_cases_test.go @@ -118,6 +118,88 @@ var enumAcceptanceTestCases = []acceptanceTestCase{ // as a validation error. In the future, we can identify this in the actual plan generation stage. expectedPlanErrorContains: errValidatingPlan.Error(), }, + { + // Exercises the CREATE TYPE ... AS ENUM path with labels containing single + // quotes. Verifies the generated DDL is correctly escaped, executes against + // Postgres, and stores the labels verbatim. + name: "create enum with single-quote labels", + oldSchemaDDL: []string{ + ` + CREATE TABLE foo(); + `, + }, + newSchemaDDL: []string{ + ` + CREATE TYPE quote_enum AS ENUM ('O''Brien', 'it''s', 'plain'); + CREATE TABLE foo( + val quote_enum DEFAULT 'plain' + ); + `, + }, + }, + { + // Exercises the ALTER TYPE ... ADD VALUE path (value appended at the end) + // with a label containing a single quote. + name: "add value with single-quote label", + oldSchemaDDL: []string{ + ` + CREATE TYPE quote_enum AS ENUM ('a', 'b'); + CREATE TABLE foo( + val quote_enum + ); + `, + }, + newSchemaDDL: []string{ + ` + CREATE TYPE quote_enum AS ENUM ('a', 'b', 'won''t'); + CREATE TABLE foo( + val quote_enum + ); + `, + }, + }, + { + // Exercises the ALTER TYPE ... ADD VALUE ... BEFORE path. The new value is + // inserted in the middle, so a BEFORE clause is emitted referencing a + // following label that contains a single quote. + name: "add value before a single-quote label", + oldSchemaDDL: []string{ + ` + CREATE TYPE quote_enum AS ENUM ('a', 'it''s'); + CREATE TABLE foo( + val quote_enum + ); + `, + }, + newSchemaDDL: []string{ + ` + CREATE TYPE quote_enum AS ENUM ('a', 'mid', 'it''s'); + CREATE TABLE foo( + val quote_enum + ); + `, + }, + }, + { + // Regression test for the enum-label SQL injection: a label crafted as an + // injection payload must be stored inertly, not executed. If escaping ever + // regresses, applying the generated plan would attempt to DROP the table and + // the pg_dump comparison would fail. + name: "enum label with sql injection payload", + oldSchemaDDL: []string{ + ` + CREATE TABLE foo(); + `, + }, + newSchemaDDL: []string{ + ` + CREATE TYPE evil AS ENUM ('x''); DROP TABLE foo; --'); + CREATE TABLE foo( + val evil + ); + `, + }, + }, } func TestEnumTestCases(t *testing.T) { diff --git a/internal/schema/escape_test.go b/internal/schema/escape_test.go new file mode 100644 index 0000000..e3b3086 --- /dev/null +++ b/internal/schema/escape_test.go @@ -0,0 +1,60 @@ +package schema + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestEscapeLiteral(t *testing.T) { + for _, tc := range []struct { + name string + input string + expected string + }{ + { + name: "simple value", + input: "red", + expected: "'red'", + }, + { + name: "value with single quote", + input: "it's", + expected: "'it''s'", + }, + { + name: "value with multiple single quotes", + input: "it's a 'test'", + expected: "'it''s a ''test'''", + }, + { + name: "empty string", + input: "", + expected: "''", + }, + { + name: "value with backslash", + input: `back\slash`, + expected: `'back\slash'`, + }, + { + name: "value with null byte stripped", + input: "null\x00byte", + expected: "'nullbyte'", + }, + { + name: "SQL injection attempt via single quote", + input: "x'); DROP TABLE foo; --", + expected: "'x''); DROP TABLE foo; --'", + }, + { + name: "COPY TO PROGRAM injection attempt", + input: "x'); COPY(SELECT 1)TO PROGRAM 'id>/tmp/pwned'; --", + expected: "'x''); COPY(SELECT 1)TO PROGRAM ''id>/tmp/pwned''; --'", + }, + } { + t.Run(tc.name, func(t *testing.T) { + assert.Equal(t, tc.expected, EscapeLiteral(tc.input)) + }) + } +} diff --git a/internal/schema/schema.go b/internal/schema/schema.go index b3945d5..527ddc8 100644 --- a/internal/schema/schema.go +++ b/internal/schema/schema.go @@ -1608,6 +1608,14 @@ func EscapeIdentifier(name string) string { return pgx.Identifier{name}.Sanitize() } +// EscapeLiteral returns a safely escaped SQL string literal, enclosed in single +// quotes. Single quotes within the value are doubled per the SQL standard, and +// null bytes are stripped as they are not valid in PostgreSQL string literals. +func EscapeLiteral(val string) string { + val = strings.ReplaceAll(val, string([]byte{0}), "") + return "'" + strings.ReplaceAll(val, "'", "''") + "'" +} + // relOptionsToMap converts pg_catalog.pg_class.reloptions to a map. func relOptionsToMap(vals []string) (map[string]string, error) { out := make(map[string]string) diff --git a/pkg/diff/enum_sql_generator.go b/pkg/diff/enum_sql_generator.go index de5ab4f..bac74f6 100644 --- a/pkg/diff/enum_sql_generator.go +++ b/pkg/diff/enum_sql_generator.go @@ -17,7 +17,7 @@ type enumSQLGenerator struct{} func (e *enumSQLGenerator) Add(enum schema.Enum) ([]Statement, error) { var escapedEnumVals []string for _, val := range enum.Labels { - escapedEnumVals = append(escapedEnumVals, fmt.Sprintf("'%s'", val)) + escapedEnumVals = append(escapedEnumVals, schema.EscapeLiteral(val)) } return []Statement{ { @@ -76,9 +76,9 @@ func (e *enumSQLGenerator) Alter(diff enumDiff) ([]Statement, error) { continue } sb := strings.Builder{} - sb.WriteString(fmt.Sprintf("ALTER TYPE %s ADD VALUE '%s'", diff.new.GetFQEscapedName(), val)) + sb.WriteString(fmt.Sprintf("ALTER TYPE %s ADD VALUE %s", diff.new.GetFQEscapedName(), schema.EscapeLiteral(val))) if i < len(diff.new.Labels)-1 { - sb.WriteString(fmt.Sprintf(" BEFORE '%s'", diff.new.Labels[i+1])) + sb.WriteString(fmt.Sprintf(" BEFORE %s", schema.EscapeLiteral(diff.new.Labels[i+1]))) } stmts = append(stmts, Statement{ DDL: sb.String(), diff --git a/pkg/diff/enum_sql_generator_test.go b/pkg/diff/enum_sql_generator_test.go new file mode 100644 index 0000000..d97a03c --- /dev/null +++ b/pkg/diff/enum_sql_generator_test.go @@ -0,0 +1,105 @@ +package diff + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/stripe/pg-schema-diff/internal/schema" +) + +func TestEnumSQLGenerator_Add_EscapesLabels(t *testing.T) { + gen := &enumSQLGenerator{} + + for _, tc := range []struct { + name string + labels []string + expectedDDL string + }{ + { + name: "simple labels", + labels: []string{"red", "green", "blue"}, + expectedDDL: `CREATE TYPE "public"."color" AS ENUM ('red', 'green', 'blue')`, + }, + { + name: "label with single quote", + labels: []string{"it's", "fine"}, + expectedDDL: `CREATE TYPE "public"."color" AS ENUM ('it''s', 'fine')`, + }, + { + name: "SQL injection attempt", + labels: []string{"x'); DROP TABLE foo; --"}, + expectedDDL: `CREATE TYPE "public"."color" AS ENUM ('x''); DROP TABLE foo; --')`, + }, + { + name: "COPY TO PROGRAM injection attempt", + labels: []string{"x'); COPY(SELECT 1)TO PROGRAM 'id'; --"}, + expectedDDL: `CREATE TYPE "public"."color" AS ENUM ('x''); COPY(SELECT 1)TO PROGRAM ''id''; --')`, + }, + } { + t.Run(tc.name, func(t *testing.T) { + stmts, err := gen.Add(schema.Enum{ + SchemaQualifiedName: schema.SchemaQualifiedName{ + SchemaName: "public", + EscapedName: `"color"`, + }, + Labels: tc.labels, + }) + require.NoError(t, err) + require.Len(t, stmts, 1) + assert.Equal(t, tc.expectedDDL, stmts[0].DDL) + }) + } +} + +func TestEnumSQLGenerator_Alter_EscapesLabels(t *testing.T) { + gen := &enumSQLGenerator{} + + stmts, err := gen.Alter(enumDiff{ + oldAndNew: oldAndNew[schema.Enum]{ + old: schema.Enum{ + SchemaQualifiedName: schema.SchemaQualifiedName{ + SchemaName: "public", + EscapedName: `"status"`, + }, + Labels: []string{"active"}, + }, + new: schema.Enum{ + SchemaQualifiedName: schema.SchemaQualifiedName{ + SchemaName: "public", + EscapedName: `"status"`, + }, + Labels: []string{"active", "it's complicated"}, + }, + }, + }) + require.NoError(t, err) + require.Len(t, stmts, 1) + assert.Equal(t, `ALTER TYPE "public"."status" ADD VALUE 'it''s complicated'`, stmts[0].DDL) +} + +func TestEnumSQLGenerator_Alter_EscapesBEFORE(t *testing.T) { + gen := &enumSQLGenerator{} + + stmts, err := gen.Alter(enumDiff{ + oldAndNew: oldAndNew[schema.Enum]{ + old: schema.Enum{ + SchemaQualifiedName: schema.SchemaQualifiedName{ + SchemaName: "public", + EscapedName: `"status"`, + }, + Labels: []string{"active", "it's complicated"}, + }, + new: schema.Enum{ + SchemaQualifiedName: schema.SchemaQualifiedName{ + SchemaName: "public", + EscapedName: `"status"`, + }, + Labels: []string{"active", "new'val", "it's complicated"}, + }, + }, + }) + require.NoError(t, err) + require.Len(t, stmts, 1) + assert.Equal(t, `ALTER TYPE "public"."status" ADD VALUE 'new''val' BEFORE 'it''s complicated'`, stmts[0].DDL) +}