Escape collate (#220) (#554)

JetBrains · Apr 30, 2019 · 4c7a9a4 · 4c7a9a4
1 parent 96e78f9
commit 4c7a9a4
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 11 deletions.
diff --git a/exposed/src/main/kotlin/org/jetbrains/exposed/sql/ColumnType.kt b/exposed/src/main/kotlin/org/jetbrains/exposed/sql/ColumnType.kt
@@ -265,17 +265,13 @@ class DateColumnType(val time: Boolean): ColumnType() {
 }
 
 abstract class StringColumnType(val collate: String? = null) : ColumnType() {
-    private val charactersToEscape = mapOf(
-            '\'' to "\'\'",
-//            '\"' to "\"\"", // no need to escape double quote as we put string in single quotes
-            '\r' to "\\r",
-            '\n' to "\\n")
+    protected fun escape(value: String): String {
+        return value.map { charactersToEscape[it] ?: it }.joinToString("")
+    }
 
     override fun nonNullValueToString(value: Any): String = buildString {
         append('\'')
-        value.toString().forEach {
-            append(charactersToEscape[it] ?: it)
-        }
+        append(escape(value.toString()))
         append('\'')
     }
 
@@ -284,14 +280,23 @@ abstract class StringColumnType(val collate: String? = null) : ColumnType() {
         is ByteArray -> String(value)
         else -> value
     }
+
+    companion object {
+        private val charactersToEscape = mapOf(
+                '\'' to "\'\'",
+//            '\"' to "\"\"", // no need to escape double quote as we put string in single quotes
+                '\r' to "\\r",
+                '\n' to "\\n"
+        )
+    }
 }
 
 open class VarCharColumnType(val colLength: Int = 255, collate: String? = null) : StringColumnType(collate)  {
     override fun sqlType(): String = buildString {
         append("VARCHAR($colLength)")
 
         if (collate != null) {
-            append(" COLLATE $collate")
+            append(" COLLATE ${escape(collate)}")
         }
     }
 }
@@ -301,7 +306,7 @@ open class TextColumnType(collate: String? = null) : StringColumnType(collate) {
         append(currentDialect.dataTypeProvider.textType())
 
         if (collate != null) {
-            append(" COLLATE $collate")
+            append(" COLLATE ${escape(collate)}")
         }
     }
 }

diff --git a/exposed/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/DDLTests.kt b/exposed/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/DDLTests.kt
@@ -2,7 +2,6 @@ package org.jetbrains.exposed.sql.tests.shared
 
 import org.jetbrains.exposed.dao.*
 import org.jetbrains.exposed.sql.*
-import org.jetbrains.exposed.sql.SqlExpressionBuilder.eq
 import org.jetbrains.exposed.sql.tests.DatabaseTestsBase
 import org.jetbrains.exposed.sql.tests.TestDB
 import org.jetbrains.exposed.sql.transactions.TransactionManager
@@ -15,6 +14,7 @@ import org.postgresql.util.PGobject
 import java.sql.SQLException
 import java.util.*
 import javax.sql.rowset.serial.SerialBlob
+import kotlin.test.assertEquals
 import kotlin.test.assertFalse
 import kotlin.test.assertNotNull
 
@@ -510,6 +510,20 @@ class DDLTests : DatabaseTestsBase() {
         }
     }
 
+    @Test fun testEscapeStringColumnType() {
+        withDb(TestDB.H2) {
+            assertEquals("VARCHAR(255) COLLATE utf8_general_ci", VarCharColumnType(collate = "utf8_general_ci").sqlType())
+            assertEquals("VARCHAR(255) COLLATE injected''code", VarCharColumnType(collate = "injected'code").sqlType())
+            assertEquals("'value'", VarCharColumnType().nonNullValueToString("value"))
+            assertEquals("'injected''value'", VarCharColumnType().nonNullValueToString("injected'value"))
+
+            assertEquals("TEXT COLLATE utf8_general_ci", TextColumnType(collate = "utf8_general_ci").sqlType())
+            assertEquals("TEXT COLLATE injected''code", TextColumnType(collate = "injected'code").sqlType())
+            assertEquals("'value'", TextColumnType().nonNullValueToString("value"))
+            assertEquals("'injected''value'", TextColumnType().nonNullValueToString("injected'value"))
+        }
+    }
+
     @Test fun addAutoPrimaryKey() {
         val tableName = "Foo"
         val initialTable = object : Table(tableName) {