fix(c/driver/postgresql): Prevent SQL Injection in GetTableSchema (#657)

Fixes #643.
apache · May 9, 2023 · e0135af · e0135af
1 parent 5b01ed2
commit e0135af
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 1 deletion.
diff --git a/c/driver/postgresql/connection.cc b/c/driver/postgresql/connection.cc
@@ -105,7 +105,7 @@ AdbcStatusCode PostgresConnection::GetTableSchema(const char* catalog,
     return ADBC_STATUS_INVALID_ARGUMENT;
   }
 
-  int ret = StringBuilderAppend(&query, "%s%s", table_name, "'::regclass::oid");
+  int ret = StringBuilderAppend(&query, "%s%s", table, "'::regclass::oid");
   PQfreemem(table);
 
   if (ret != 0) return ADBC_STATUS_INTERNAL;

diff --git a/c/driver/postgresql/postgresql_test.cc b/c/driver/postgresql/postgresql_test.cc
@@ -30,6 +30,7 @@
 #include "validation/adbc_validation_util.h"
 
 using adbc_validation::IsOkStatus;
+using adbc_validation::IsStatus;
 
 class PostgresQuirks : public adbc_validation::DriverQuirks {
  public:
@@ -97,6 +98,41 @@ class PostgresConnectionTest : public ::testing::Test,
  protected:
   PostgresQuirks quirks_;
 };
+
+TEST_F(PostgresConnectionTest, MetadataGetTableSchemaInjection) {
+  if (!quirks()->supports_bulk_ingest()) {
+    GTEST_SKIP();
+  }
+  ASSERT_THAT(AdbcConnectionNew(&connection, &error), IsOkStatus(&error));
+  ASSERT_THAT(AdbcConnectionInit(&connection, &database, &error), IsOkStatus(&error));
+  ASSERT_THAT(quirks()->DropTable(&connection, "bulk_ingest", &error),
+              IsOkStatus(&error));
+  ASSERT_THAT(quirks()->EnsureSampleTable(&connection, "bulk_ingest", &error),
+              IsOkStatus(&error));
+
+  adbc_validation::Handle<ArrowSchema> schema;
+  ASSERT_THAT(AdbcConnectionGetTableSchema(&connection, /*catalog=*/nullptr,
+                                           /*db_schema=*/nullptr,
+                                           "0'::int; DROP TABLE bulk_ingest;--",
+                                           &schema.value, &error),
+              IsStatus(ADBC_STATUS_IO, &error));
+
+  ASSERT_THAT(
+      AdbcConnectionGetTableSchema(&connection, /*catalog=*/nullptr,
+                                   /*db_schema=*/"0'::int; DROP TABLE bulk_ingest;--",
+                                   "DROP TABLE bulk_ingest;", &schema.value, &error),
+      IsStatus(ADBC_STATUS_IO, &error));
+
+  ASSERT_THAT(AdbcConnectionGetTableSchema(&connection, /*catalog=*/nullptr,
+                                           /*db_schema=*/nullptr, "bulk_ingest",
+                                           &schema.value, &error),
+              IsOkStatus(&error));
+
+  ASSERT_NO_FATAL_FAILURE(adbc_validation::CompareSchema(
+      &schema.value, {{"int64s", NANOARROW_TYPE_INT64, true},
+                      {"strings", NANOARROW_TYPE_STRING, true}}));
+}
+
 ADBCV_TEST_CONNECTION(PostgresConnectionTest)
 
 class PostgresStatementTest : public ::testing::Test,