fix: escape enum labels in generated SQL to prevent SQL injection

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) {

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))
+		})
+	}
+}

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)

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(),

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)
+}