Merge pull request #40 from Steveiwonder/bug/34

fixes #34 - Ensure columns / parameter names are correctly escaped
Steveiwonder · Oct 11, 2023 · 2056344 · 2056344
2 parents aad2b4c + a8e9ce3
commit 2056344
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 11 deletions.
diff --git a/README.md b/README.md
@@ -164,7 +164,7 @@ When using data type Sql this allows you to get values from other tables within
 
 | Property Name | Values |
 | ------------- | ------ |
-| query | The query to use for the lookup, the current row will be passed into the query as parameters for use, see the example config above that uses @UserId |
+| query | The query to use for the lookup, the current row will be passed into the query as parameters for use, see the example config above that uses @UserId. Columns are passed back into the query as parameters, any columns with spaces in their name, will be replaced with '_' . *e.g. A column name "User Id" would become "User_Id"* |
 | valueHandling | "Null" or "KeepValue". If the query executes and no data is returned, this tells the masker what to do, null  will set the value to Null while KeepValue will keep the existing value on that row |
 
 ## Data types

diff --git a/src/DataMasker/DataSources/SqlDataSource.cs b/src/DataMasker/DataSources/SqlDataSource.cs
@@ -70,7 +70,7 @@ public void UpdateRow(
             using (SqlConnection connection = new SqlConnection(_connectionString))
             {
                 connection.Open();
-                connection.Execute(BuildUpdateSql(tableConfig), row, null, commandType: CommandType.Text);
+                connection.Execute(BuildUpdateSql(tableConfig), Utils.Utils.MakeParamNamesSafe(row), null, commandType: CommandType.Text);
             }
         }
 
@@ -105,7 +105,8 @@ public void UpdateRows(
 
 
                     string sql = BuildUpdateSql(config);
-                    connection.Execute(sql, batch.Items, sqlTransaction, null, CommandType.Text);
+                    var safeDictionaries = batch.Items.Select(Utils.Utils.MakeParamNamesSafe);
+                    connection.Execute(sql, safeDictionaries, sqlTransaction, null, CommandType.Text);
 
                     if (_sourceConfig.DryRun)
                     {
@@ -147,7 +148,7 @@ private string BuildUpdateSql(
             string sql = $"UPDATE [{tableConfig.Schema}].[{tableConfig.Name}] SET ";
 
             sql += tableConfig.Columns.GetUpdateColumns();
-            sql += $" WHERE [{tableConfig.PrimaryKeyColumn}] = @{tableConfig.PrimaryKeyColumn}";
+            sql += $" WHERE {Utils.Utils.MakeColumnNameSafe(tableConfig.PrimaryKeyColumn)} = @{Utils.Utils.MakeParamNameSafe(tableConfig.PrimaryKeyColumn)}";
             return sql;
         }
 

diff --git a/src/DataMasker/Extensions.cs b/src/DataMasker/Extensions.cs
@@ -19,7 +19,7 @@ public static string GetSelectColumns(
             this IList<ColumnConfig> columns,
             string primaryKeyColumn)
         {
-            IList<string> columnNames = new List<string>(columns.Select(x => $"[{x.Name}]"));
+            IList<string> columnNames = new List<string>(columns.Select(x => $"{Utils.Utils.MakeColumnNameSafe(x.Name)}"));
             columnNames.Insert(0, primaryKeyColumn);
             return string.Join(", ", columnNames);
         }
@@ -37,7 +37,7 @@ public static string GetUpdateColumns(
             return string.Join(
                                ", ",
                                columns.Where(x => !x.Ignore)
-                                      .Select(x => $"[{x.Name}] = @{paramPrefix}{x.Name}"));
+                                      .Select(x => $"{Utils.Utils.MakeColumnNameSafe(x.Name)} = @{paramPrefix}{Utils.Utils.MakeParamNameSafe(x.Name)}"));
         }
     }
 }
diff --git a/src/DataMasker/SqlDataPovider.cs b/src/DataMasker/SqlDataPovider.cs
@@ -3,8 +3,8 @@
 using DataMasker.Models;
 using System.Data.SqlClient;
 using Dapper;
-using System.Collections.Generic;
-
+using System.Collections.Generic;
+
 namespace DataMasker
 {
 
@@ -22,8 +22,10 @@ public bool CanProvide(DataType dataType)
     }
 
     public object GetValue(ColumnConfig columnConfig, IDictionary<string, object> obj, Name.Gender? gender)
-    {
-      DynamicParameters dynamicParameters = new DynamicParameters(obj);
+    {
+
+      IDictionary<string, object> safeObj = Utils.Utils.MakeParamNamesSafe(obj);
+      DynamicParameters dynamicParameters = new DynamicParameters(safeObj);
       object newValue = _connection.ExecuteScalar(columnConfig.SqlValue.Query, dynamicParameters);
       if (newValue == null && columnConfig.SqlValue.ValueHandling == NotFoundValueHandling.KeepValue)
       {

diff --git a/src/DataMasker/Utils/Utils.cs b/src/DataMasker/Utils/Utils.cs
@@ -1,5 +1,7 @@
 using Bogus.DataSets;
-
+using System.Collections.Generic;
+using System.Linq;
+
 namespace DataMasker.Utils
 {
     public static class Utils
@@ -33,5 +35,30 @@ public static class Utils
 
             return null;
         }
+
+        internal static IDictionary<string, object> MakeParamNamesSafe(IDictionary<string, object> param)
+        {
+            return param.ToDictionary(d => $"{MakeParamNameSafe(d.Key)}", d => d.Value);
+        }
+
+        internal static string MakeParamNameSafe(string paramName)
+        {
+            return paramName.Replace(" ", "_");
+        }
+
+        internal static string MakeColumnNameSafe(string paramName)
+        {
+            if (paramName[0] != '[')
+            {
+                paramName = $"[{paramName}";
+            }
+
+            if (paramName[paramName.Length-1] != ']')
+            {
+                paramName = $"{paramName}]";
+            }
+
+            return paramName;
+        }
     }
 }