Allow parametrized generalization expressions

Author: pdobacz

pg_diffix/query/anonymization.h

@@ -16,7 +16,7 @@ typedef struct AnonQueryLinks AnonQueryLinks;
  * Transforms subqueries accessing personal relations into anonymizing subqueries.
  * Returned data is used during plan rewrite.
  */
-extern AnonQueryLinks *compile_query(Query *query, List *personal_relations);
+extern AnonQueryLinks *compile_query(Query *query, List *personal_relations, ParamListInfo bound_params);
 
 /*
  * Calls `rewrite_plan` for each item in a list of Plan nodes.

pg_diffix/query/validation.h

@@ -27,7 +27,7 @@ extern void verify_anonymization_requirements(Query *query);
  * Verifies restrictions on bucket expressions, operates on an anonymizing query.
  * If requirements are not met, an error is reported and execution is halted.
  */
-extern void verify_bucket_expressions(Query *query);
+extern void verify_bucket_expressions(Query *query, ParamListInfo bound_params);
 
 /*
  * Returns `true` if the given list of `RangeTblEntry` from `ExecutorCheckPerms` does not access `pg_catalog`

pg_diffix/utils.h

@@ -1,7 +1,9 @@
 #ifndef PG_DIFFIX_UTILS_H
 #define PG_DIFFIX_UTILS_H
 
+#include "nodes/params.h"
 #include "nodes/pg_list.h"
+#include "nodes/primnodes.h"
 #include "utils/datum.h"
 
 /*-------------------------------------------------------------------------
@@ -119,4 +121,68 @@ static inline List *hash_set_union(List *dst_set, const List *src_set)
 
 #endif
 
+/*-------------------------------------------------------------------------
+ * Node utils
+ *-------------------------------------------------------------------------
+ */
+
+static inline ParamExternData *get_param_data(ParamListInfo bound_params, int one_based_paramid)
+{
+#if PG_MAJORVERSION_NUM == 13
+  int paramid = one_based_paramid;
+#else
+  int paramid = one_based_paramid - 1;
+#endif
+  if (bound_params->paramFetch != NULL)
+    return bound_params->paramFetch(bound_params, paramid, true, NULL);
+  else
+    return &bound_params->params[paramid];
+}
+
+static inline bool is_simple_constant(Node *node)
+{
+  return IsA(node, Const) || (IsA(node, Param) && ((Param *)node)->paramkind == PARAM_EXTERN);
+}
+
+static inline void get_simple_constant_typed_value(Node *node, ParamListInfo bound_params, Oid *type, Datum *value, bool *isnull)
+{
+  if (IsA(node, Const))
+  {
+    Const *const_expr = (Const *)node;
+    *type = const_expr->consttype;
+    *value = const_expr->constvalue;
+    *isnull = const_expr->constisnull;
+  }
+  else if (IsA(node, Param) && ((Param *)node)->paramkind == PARAM_EXTERN)
+  {
+    Param *param_expr = (Param *)node;
+    ParamExternData *param_data = get_param_data(bound_params, param_expr->paramid);
+    *type = param_data->ptype;
+    *value = param_data->value;
+    *isnull = param_data->isnull;
+  }
+  else
+  {
+    FAILWITH("Attempted to get simple constant value of non-Const, non-PARAM_EXTERN node");
+  }
+}
+
+static inline int get_simple_constant_location(Node *node)
+{
+  if (IsA(node, Const))
+  {
+    Const *const_expr = (Const *)node;
+    return const_expr->location;
+  }
+  else if (IsA(node, Param) && ((Param *)node)->paramkind == PARAM_EXTERN)
+  {
+    Param *param_expr = (Param *)node;
+    return param_expr->location;
+  }
+  else
+  {
+    FAILWITH("Attempted to get simple constant value of non-Const, non-PARAM_EXTERN node");
+  }
+}
+
 #endif /* PG_DIFFIX_UTILS_H */

src/hooks.c

@@ -47,7 +47,7 @@ static void pg_diffix_post_parse_analyze(ParseState *pstate, Query *query, Jumbl
 }
 #endif
 
-static AnonQueryLinks *prepare_query(Query *query)
+static AnonQueryLinks *prepare_query(Query *query, ParamListInfo bound_params)
 {
   /* Do nothing for sessions with direct access. */
   if (get_session_access_level() == ACCESS_DIRECT)
@@ -70,7 +70,7 @@ static AnonQueryLinks *prepare_query(Query *query)
    */
   config_validate();
 
-  AnonQueryLinks *links = compile_query(query, personal_relations);
+  AnonQueryLinks *links = compile_query(query, personal_relations, bound_params);
 
   DEBUG_LOG("Compiled query (Query ID=%lu) (User ID=%u) %s", query->queryId, GetSessionUserId(), nodeToString(query));
 
@@ -88,7 +88,7 @@ static PlannedStmt *pg_diffix_planner(
 
   DEBUG_LOG("Statement (Query ID=%lu) (User ID=%u): %s", query->queryId, GetSessionUserId(), query_string);
 
-  AnonQueryLinks *links = prepare_query(query);
+  AnonQueryLinks *links = prepare_query(query, boundParams);
 
   planner_hook_type planner = (prev_planner_hook ? prev_planner_hook : standard_planner);
   PlannedStmt *plan = planner(query, query_string, cursorOptions, boundParams);

src/query/anonymization.c

@@ -316,6 +316,7 @@ static void append_seed_material(
 typedef struct CollectMaterialContext
 {
   Query *query;
+  ParamListInfo bound_params;
   char material[MAX_SEED_MATERIAL_SIZE];
 } CollectMaterialContext;
 
@@ -356,14 +357,17 @@ static bool collect_seed_material(Node *node, CollectMaterialContext *context)
     append_seed_material(context->material, attribute_name, '.');
   }
 
-  if (IsA(node, Const))
+  if (is_simple_constant(node))
   {
-    Const *const_expr = (Const *)node;
+    Oid type;
+    Datum value;
+    bool isnull;
+    get_simple_constant_typed_value(node, context->bound_params, &type, &value, &isnull);
 
-    if (!is_supported_numeric_type(const_expr->consttype))
-      FAILWITH_LOCATION(const_expr->location, "Unsupported constant type used in bucket definition!");
+    if (!is_supported_numeric_type(type))
+      FAILWITH_LOCATION(get_simple_constant_location(node), "Unsupported constant type used in bucket definition!");
 
-    double const_as_double = numeric_value_to_double(const_expr->consttype, const_expr->constvalue);
+    double const_as_double = numeric_value_to_double(type, value);
     char const_as_string[DOUBLE_SHORTEST_DECIMAL_LEN];
     double_to_shortest_decimal_buf(const_as_double, const_as_string);
     append_seed_material(context->material, const_as_string, ',');
@@ -377,7 +381,7 @@ static bool collect_seed_material(Node *node, CollectMaterialContext *context)
  * Computes the SQL part of the bucket seed by combining the unique grouping expressions' seed material hashes.
  * Grouping clause (if any) must be made explicit before calling this.
  */
-static seed_t prepare_bucket_seeds(Query *query)
+static seed_t prepare_bucket_seeds(Query *query, ParamListInfo bound_params)
 {
   List *seed_material_hash_set = NULL;
 
@@ -388,7 +392,7 @@ static seed_t prepare_bucket_seeds(Query *query)
     Node *expr = lfirst(cell);
 
     /* Start from empty string and append material pieces for each non-cast expression. */
-    CollectMaterialContext collect_context = {.query = query, .material = ""};
+    CollectMaterialContext collect_context = {.query = query, .bound_params = bound_params, .material = ""};
     collect_seed_material(expr, &collect_context);
 
     /* Keep materials with unique hashes to avoid them cancelling each other. */
@@ -578,15 +582,15 @@ static void wrap_having_qual(Query *query)
       COERCE_EXPLICIT_CALL);
 }
 
-static void compile_anonymizing_query(Query *query, List *personal_relations, AnonQueryLinks *anon_links)
+static void compile_anonymizing_query(Query *query, List *personal_relations, AnonQueryLinks *anon_links, ParamListInfo bound_params)
 {
   verify_anonymization_requirements(query);
 
   AnonymizationContext *anon_context = make_query_anonymizing(query, personal_relations);
 
-  verify_bucket_expressions(query);
+  verify_bucket_expressions(query, bound_params);
 
-  anon_context->sql_seed = prepare_bucket_seeds(query);
+  anon_context->sql_seed = prepare_bucket_seeds(query, bound_params);
 
   link_anon_context(query, anon_links, anon_context);
 
@@ -613,6 +617,7 @@ typedef struct QueryCompileContext
 {
   List *personal_relations;
   AnonQueryLinks *anon_links;
+  ParamListInfo bound_params;
 } QueryCompileContext;
 
 static bool compile_query_walker(Node *node, QueryCompileContext *context)
@@ -651,19 +656,20 @@ static bool compile_query_walker(Node *node, QueryCompileContext *context)
   {
     Query *query = (Query *)node;
     if (is_anonymizing_query(query, context->personal_relations))
-      compile_anonymizing_query(query, context->personal_relations, context->anon_links);
+      compile_anonymizing_query(query, context->personal_relations, context->anon_links, context->bound_params);
     else
       query_tree_walker(query, compile_query_walker, context, QTW_EXAMINE_RTES_AFTER);
   }
 
   return expression_tree_walker(node, compile_query_walker, context);
 }
 
-AnonQueryLinks *compile_query(Query *query, List *personal_relations)
+AnonQueryLinks *compile_query(Query *query, List *personal_relations, ParamListInfo bound_params)
 {
   QueryCompileContext context = {
       .personal_relations = personal_relations,
       .anon_links = palloc0(sizeof(AnonQueryLinks)),
+      .bound_params = bound_params,
   };
 
   compile_query_walker((Node *)query, &context);

src/query/validation.c

@@ -54,6 +54,8 @@ void verify_utility_command(Node *utility_stmt)
     case T_DeallocateStmt:
     case T_FetchStmt:
     case T_ClosePortalStmt:
+    case T_PrepareStmt:
+    case T_ExecuteStmt:
       break;
     default:
       FAILWITH("Statement requires direct access level.");
@@ -225,7 +227,7 @@ static void verify_bucket_expression(Node *node)
 
     for (int i = 1; i < list_length(func_expr->args); i++)
     {
-      if (!IsA(unwrap_cast((Node *)list_nth(func_expr->args, i)), Const))
+      if (!is_simple_constant(unwrap_cast((Node *)list_nth(func_expr->args, i))))
         FAILWITH_LOCATION(func_expr->location, "Non-primary arguments for a bucket function have to be simple constants.");
     }
   }
@@ -234,10 +236,9 @@ static void verify_bucket_expression(Node *node)
     OpExpr *op_expr = (OpExpr *)node;
     FAILWITH_LOCATION(op_expr->location, "Use of operators to define buckets is not supported.");
   }
-  else if (IsA(node, Const))
+  else if (is_simple_constant(node))
   {
-    Const *const_expr = (Const *)node;
-    FAILWITH_LOCATION(const_expr->location, "Simple constants are not allowed as bucket expressions.");
+    FAILWITH_LOCATION(get_simple_constant_location(node), "Simple constants are not allowed as bucket expressions.");
   }
   else if (IsA(node, RelabelType))
   {
@@ -262,14 +263,17 @@ static void verify_bucket_expression(Node *node)
   }
 }
 
-static void verify_substring(FuncExpr *func_expr)
+static void verify_substring(FuncExpr *func_expr, ParamListInfo bound_params)
 {
   Node *node = unwrap_cast(list_nth(func_expr->args, 1));
-  Assert(IsA(node, Const)); /* Checked by prior validations */
-  Const *second_arg = (Const *)node;
-
-  if (DatumGetUInt32(second_arg->constvalue) != 1)
-    FAILWITH_LOCATION(second_arg->location, "Generalization used in the query is not allowed in untrusted access level.");
+  Assert(is_simple_constant(node)); /* Checked by prior validations */
+  Oid type;
+  Datum value;
+  bool isnull;
+  get_simple_constant_typed_value(node, bound_params, &type, &value, &isnull);
+
+  if (DatumGetUInt32(value) != 1)
+    FAILWITH_LOCATION(get_simple_constant_location(node), "Generalization used in the query is not allowed in untrusted access level.");
 }
 
 /* money-style numbers, i.e. 1, 2, or 5 preceeded by or followed by zeros: ⟨... 0.1, 0.2, 0.5, 1, 2, 5, 10, ...⟩ */
@@ -286,29 +290,32 @@ static bool is_money_style(double number)
 }
 
 /* Expects the expression being the second argument to `round_by` et al. */
-static void verify_bin_size(Node *range_expr)
+static void verify_bin_size(Node *range_expr, ParamListInfo bound_params)
 {
   Node *range_node = unwrap_cast(range_expr);
-  Assert(IsA(range_node, Const)); /* Checked by prior validations */
-  Const *range_const = (Const *)range_node;
+  Assert(is_simple_constant(range_node)); /* Checked by prior validations */
+  Oid type;
+  Datum value;
+  bool isnull;
+  get_simple_constant_typed_value(range_node, bound_params, &type, &value, &isnull);
 
-  if (!is_supported_numeric_type(range_const->consttype))
-    FAILWITH_LOCATION(range_const->location, "Unsupported constant type used in generalization.");
+  if (!is_supported_numeric_type(type))
+    FAILWITH_LOCATION(get_simple_constant_location(range_node), "Unsupported constant type used in generalization.");
 
-  if (!is_money_style(numeric_value_to_double(range_const->consttype, range_const->constvalue)))
-    FAILWITH_LOCATION(range_const->location, "Generalization used in the query is not allowed in untrusted access level.");
+  if (!is_money_style(numeric_value_to_double(type, value)))
+    FAILWITH_LOCATION(get_simple_constant_location(range_node), "Generalization used in the query is not allowed in untrusted access level.");
 }
 
-static void verify_generalization(Node *node)
+static void verify_generalization(Node *node, ParamListInfo bound_params)
 {
   if (IsA(node, FuncExpr))
   {
     FuncExpr *func_expr = (FuncExpr *)node;
 
     if (is_substring_builtin(func_expr->funcid))
-      verify_substring(func_expr);
+      verify_substring(func_expr, bound_params);
     else if (is_implicit_range_udf_untrusted(func_expr->funcid))
-      verify_bin_size((Node *)list_nth(func_expr->args, 1));
+      verify_bin_size((Node *)list_nth(func_expr->args, 1), bound_params);
     else if (is_implicit_range_builtin_untrusted(func_expr->funcid))
       ;
     else
@@ -317,7 +324,7 @@ static void verify_generalization(Node *node)
 }
 
 /* Should be run on anonymizing queries only. */
-void verify_bucket_expressions(Query *query)
+void verify_bucket_expressions(Query *query, ParamListInfo bound_params)
 {
   AccessLevel access_level = get_session_access_level();
 
@@ -333,7 +340,7 @@ void verify_bucket_expressions(Query *query)
     Node *expr = (Node *)lfirst(cell);
     verify_bucket_expression(expr);
     if (access_level == ACCESS_ANONYMIZED_UNTRUSTED)
-      verify_generalization(expr);
+      verify_generalization(expr, bound_params);
   }
 }
 

test/expected/validation.out

@@ -361,6 +361,9 @@ SELECT city, 'aaaa' FROM test_validation GROUP BY 1, 2;
 ERROR:  [PG_DIFFIX] Simple constants are not allowed as bucket expressions.
 LINE 1: SELECT city, 'aaaa' FROM test_validation GROUP BY 1, 2;
                      ^
+PREPARE prepared_param_as_label(text) AS SELECT city, $1 FROM test_validation GROUP BY 1, 2;
+EXECUTE prepared_param_as_label('aaaa');
+ERROR:  [PG_DIFFIX] Simple constants are not allowed as bucket expressions.
 SELECT COUNT(*) FROM test_validation GROUP BY round(floor(id));
 ERROR:  [PG_DIFFIX] Primary argument for a bucket function has to be a simple column reference.
 LINE 1: SELECT COUNT(*) FROM test_validation GROUP BY round(floor(id...
@@ -586,3 +589,20 @@ SELECT diffix.ceil_by(discount, 2) from test_validation;
 ERROR:  [PG_DIFFIX] Generalization used in the query is not allowed in untrusted access level.
 LINE 1: SELECT diffix.ceil_by(discount, 2) from test_validation;
                ^
+-- Allow prepared statements with generalization constants as params, and validate them
+PREPARE prepared_floor_by(numeric) AS SELECT diffix.floor_by(discount, $1) FROM test_validation GROUP BY 1;
+EXECUTE prepared_floor_by(2.0);
+ floor_by 
+----------
+(0 rows)
+
+EXECUTE prepared_floor_by(2.1);
+ERROR:  [PG_DIFFIX] Generalization used in the query is not allowed in untrusted access level.
+PREPARE prepared_substring(int, int) AS SELECT substring(city, $1, $2) FROM test_validation GROUP BY 1;
+EXECUTE prepared_substring(1, 2);
+ substring 
+-----------
+(0 rows)
+
+EXECUTE prepared_substring(2, 3);
+ERROR:  [PG_DIFFIX] Generalization used in the query is not allowed in untrusted access level.

test/sql/validation.sql

@@ -187,6 +187,8 @@ SELECT COUNT(*) FROM test_validation GROUP BY LENGTH(city);
 SELECT COUNT(*) FROM test_validation GROUP BY city || 'xxx';
 SELECT LENGTH(city) FROM test_validation;
 SELECT city, 'aaaa' FROM test_validation GROUP BY 1, 2;
+PREPARE prepared_param_as_label(text) AS SELECT city, $1 FROM test_validation GROUP BY 1, 2;
+EXECUTE prepared_param_as_label('aaaa');
 SELECT COUNT(*) FROM test_validation GROUP BY round(floor(id));
 SELECT COUNT(*) FROM test_validation GROUP BY floor(cast(discount AS integer));
 SELECT COUNT(*) FROM test_validation GROUP BY substr(city, 1, id);
@@ -280,3 +282,11 @@ SELECT diffix.floor_by(discount, 5000000000.1) from test_validation;
 SELECT width_bucket(discount, 2, 200, 5) from test_validation;
 SELECT ceil(discount) from test_validation;
 SELECT diffix.ceil_by(discount, 2) from test_validation;
+
+-- Allow prepared statements with generalization constants as params, and validate them
+PREPARE prepared_floor_by(numeric) AS SELECT diffix.floor_by(discount, $1) FROM test_validation GROUP BY 1;
+EXECUTE prepared_floor_by(2.0);
+EXECUTE prepared_floor_by(2.1);
+PREPARE prepared_substring(int, int) AS SELECT substring(city, $1, $2) FROM test_validation GROUP BY 1;
+EXECUTE prepared_substring(1, 2);
+EXECUTE prepared_substring(2, 3);