From 52647062fef672ee689d160782976a5a4c104bf8 Mon Sep 17 00:00:00 2001 From: James Sadler Date: Sun, 26 Jul 2026 00:31:57 +1000 Subject: [PATCH 01/10] feat(mapper): fuse JSON field equality into a value-selector needle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Encrypted JSON equality (`col -> sel = value`) is not a term comparison. In EQL v3 exact equality is *containment of a value selector*: a single keyed MAC over the path and the canonicalised value together (`QueryOp::SteVecValueSelector`). One needle, built from TWO SQL operands. This is the mapper half. It breaks the implicit 1:1 correspondence between input and output operands for the first time: - New `EqlTerm::JsonValueSelector` / `EqlTermVariant::JsonValueSelector` — the value operand of a JSON field equality. Inferred for the non-JSON side of `=`/`<>` whose other side is a genuine field ACCESS (`->`, `->>`, `jsonb_path_query_first`). A bare `col = $1` on a whole encrypted JSON column is document equality and keeps its ordinary typing. - New `JsonValueSelectors` on `TypeCheckedStatement` records the one relationship node types cannot express: which operand supplies the path for which value. The mapper holds no encryption key, so it emits the composition *input*, not the needle; the proxy fuses and encrypts. - New `RewriteJsonValueSelectorEq` rewrites the comparison to `eql_v3.jsonb_contains(col, ::eql_v3.query_json)` (`<>` negates), discarding the field access. `RewriteEqlComparisonOps` now skips these — `eql_v3.eq_term` has no unique overload for a JSON query operand. - The value operand casts to `eql_v3.query_json`, alongside the existing `JsonOrd` -> `query_integer_ord` rule (helper renamed to `json_query_operand_cast_target` to cover both). A discarded selector *placeholder* stays declared in Parse but unreferenced in the SQL. PostgreSQL permits that as long as its type is known, which is what lets input and output param numbering stay identical — so Bind, ParameterDescription and the encrypt pipeline all stay positional. Verified against PostgreSQL 17: `PREPARE p (text) AS SELECT $2::jsonb` prepares and executes; the same statement with no declared types fails with "could not determine data type of parameter $1". Mapper unit tests 97 -> 102, no regressions. Proxy-side composition (the literal and Bind paths) follows in subsequent commits. Stable-Commit-Id: q-7qfnv43maha2chq6shq55sc3m0 --- packages/eql-mapper/src/eql_mapper.rs | 1 + .../src/inference/infer_type_impls/expr.rs | 136 +++++++++++- packages/eql-mapper/src/inference/mod.rs | 37 +++- .../src/inference/unifier/eql_traits.rs | 1 + .../eql-mapper/src/inference/unifier/types.rs | 19 +- .../src/inference/unifier/unify_types.rs | 4 + .../eql-mapper/src/json_value_selector.rs | 75 +++++++ packages/eql-mapper/src/lib.rs | 133 +++++++++++- .../cast_literals_as_encrypted.rs | 6 +- .../cast_params_as_encrypted.rs | 4 +- .../src/transformation_rules/helpers.rs | 26 ++- .../src/transformation_rules/mod.rs | 2 + .../rewrite_eql_comparison_ops.rs | 19 +- .../rewrite_json_value_selector_eq.rs | 195 ++++++++++++++++++ .../eql-mapper/src/type_checked_statement.rs | 20 +- 15 files changed, 659 insertions(+), 19 deletions(-) create mode 100644 packages/eql-mapper/src/json_value_selector.rs create mode 100644 packages/eql-mapper/src/transformation_rules/rewrite_json_value_selector_eq.rs diff --git a/packages/eql-mapper/src/eql_mapper.rs b/packages/eql-mapper/src/eql_mapper.rs index 50bde65b8..b2c23ceb7 100644 --- a/packages/eql-mapper/src/eql_mapper.rs +++ b/packages/eql-mapper/src/eql_mapper.rs @@ -180,6 +180,7 @@ impl<'ast> EqlMapper<'ast> { projection, params, literals, + self.inferencer.borrow().take_json_value_selectors(), Arc::new(node_types), )) } diff --git a/packages/eql-mapper/src/inference/infer_type_impls/expr.rs b/packages/eql-mapper/src/inference/infer_type_impls/expr.rs index ea9656250..9b92b9266 100644 --- a/packages/eql-mapper/src/inference/infer_type_impls/expr.rs +++ b/packages/eql-mapper/src/inference/infer_type_impls/expr.rs @@ -4,10 +4,13 @@ use crate::{ unifier::{EqlTerm, EqlValue, TokenType, Type, Value}, InferType, TypeError, }, - EqlTrait, IdentCase, TypeInferencer, + EqlTrait, IdentCase, JsonSelectorSource, Param, TypeInferencer, }; use eql_mapper_macros::trace_infer; -use sqltk::parser::ast::{AccessExpr, Array, BinaryOperator, Expr, Ident, Subscript}; +use sqltk::parser::ast::{ + self as ast, AccessExpr, Array, BinaryOperator, Expr, FunctionArg, FunctionArgExpr, + FunctionArguments, Ident, Subscript, +}; #[trace_infer] impl<'ast> InferType<'ast, Expr> for TypeInferencer<'ast> { @@ -151,6 +154,39 @@ impl<'ast> InferType<'ast, Expr> for TypeInferencer<'ast> { false }; + // Encrypted JSON field EQUALITY (`col -> sel = value`, `<>`). + // Exact equality is not a term comparison but *value-selector + // containment*: one keyed MAC over path and value together. Type + // the value operand `EqlTerm::JsonValueSelector` and record where + // its path half comes from, so the proxy can fuse the two into a + // single needle at encryption time (see `JsonValueSelectors`). + // + // Unlike the ordering case above this requires a genuine field + // ACCESS on the JSON side — a bare `col = $1` on a whole + // encrypted JSON column is document equality and must keep its + // ordinary typing. + let handled = handled + || if matches!(op, BinaryOperator::Eq | BinaryOperator::NotEq) { + match ( + self.eql_json_field_access(left), + self.eql_json_field_access(right), + ) { + (Some((json, selector)), None) => { + self.infer_json_value_selector(json, selector, right)?; + self.unify_node_with_type(expr_val, Type::native())?; + true + } + (None, Some((json, selector))) => { + self.infer_json_value_selector(json, selector, left)?; + self.unify_node_with_type(expr_val, Type::native())?; + true + } + _ => false, + } + } else { + false + }; + if !handled { get_sql_binop_rule(op).apply_constraints(self, left, right, expr_val)?; } @@ -519,4 +555,100 @@ impl<'ast> TypeInferencer<'ast> { _ => None, } } + + /// Deconstructs an encrypted-JSON **field access** into the accessed value + /// and the expression supplying its selector: + /// + /// - `col -> sel`, `col ->> sel` + /// - `jsonb_path_query_first(col, sel)` + /// + /// Returns `None` for anything else — importantly for a bare encrypted JSON + /// column, which is a whole document, not a field of one. Equality needs the + /// selector expression itself (not just the type), because the path is one + /// half of the fused value-selector needle. + fn eql_json_field_access(&self, expr: &'ast Expr) -> Option<(EqlValue, &'ast Expr)> { + let selector = match expr { + Expr::BinaryOp { + op: BinaryOperator::Arrow | BinaryOperator::LongArrow, + right, + .. + } => &**right, + + // `jsonb_path_query_first(col, sel)` — and its already-rewritten + // `eql_v3.` spelling. The selector is the second argument. + Expr::Function(function) => match &function.args { + FunctionArguments::List(list) => match list.args.as_slice() { + [_, FunctionArg::Unnamed(FunctionArgExpr::Expr(sel))] => sel, + _ => return None, + }, + _ => return None, + }, + + _ => return None, + }; + + self.eql_json_value(expr).map(|json| (json, selector)) + } + + /// Types `value` — the value half of `col -> sel = value` — as a fused + /// value selector, and records where its path half (`selector`) comes from. + /// + /// A path that is neither a literal nor a placeholder (a column reference, a + /// function call) cannot be resolved to a needle at encryption time, so the + /// fusion is declined and the comparison falls through to ordinary typing — + /// where it will fail the capability check with a clearer error than a + /// half-built needle would produce. + fn infer_json_value_selector( + &self, + json: EqlValue, + selector: &'ast Expr, + value: &'ast Expr, + ) -> Result<(), TypeError> { + let Some(source) = Self::json_selector_source(selector) else { + return Ok(()); + }; + + self.unify_node_with_type( + value, + Type::Value(Value::Eql(EqlTerm::JsonValueSelector(json))), + )?; + + match Self::as_ast_value(value) { + Some(ast::Value::Placeholder(placeholder)) => { + if let Ok(param) = Param::try_from(placeholder) { + self.record_json_value_selector_param(param, source); + } + } + Some(node) => self.record_json_value_selector_literal(node, source), + None => {} + } + + Ok(()) + } + + /// Classifies the path half of a fused value selector: a placeholder yields + /// the param it will arrive in, a literal yields its text inline. + fn json_selector_source(selector: &'ast Expr) -> Option { + match Self::as_ast_value(selector)? { + ast::Value::Placeholder(placeholder) => Param::try_from(placeholder) + .ok() + .map(JsonSelectorSource::Param), + ast::Value::SingleQuotedString(s) + | ast::Value::DoubleQuotedString(s) + | ast::Value::EscapedStringLiteral(s) => Some(JsonSelectorSource::Literal(s.clone())), + ast::Value::Number(n, _) => Some(JsonSelectorSource::Literal(n.to_string())), + _ => None, + } + } + + /// The [`ast::Value`] an expression ultimately is, seeing through casts + /// (`$1::jsonb`, `'a'::text`). Casts are common on both halves — the client + /// may write them and earlier rules may add them. + fn as_ast_value(expr: &'ast Expr) -> Option<&'ast ast::Value> { + match expr { + Expr::Value(value_with_span) => Some(&value_with_span.value), + Expr::Cast { expr, .. } => Self::as_ast_value(expr), + _ => None, + } + } } diff --git a/packages/eql-mapper/src/inference/mod.rs b/packages/eql-mapper/src/inference/mod.rs index 05fbb1eb0..58305159c 100644 --- a/packages/eql-mapper/src/inference/mod.rs +++ b/packages/eql-mapper/src/inference/mod.rs @@ -18,7 +18,9 @@ use sqltk::parser::ast::{ }; use sqltk::{into_control_flow, AsNodeKey, Break, Visitable, Visitor}; -use crate::{ScopeError, ScopeTracker, TableResolver}; +use crate::{ + JsonSelectorSource, JsonValueSelectors, Param, ScopeError, ScopeTracker, TableResolver, +}; pub(crate) use registry::*; pub(crate) use sequence::*; @@ -51,6 +53,12 @@ pub struct TypeInferencer<'ast> { /// Implements the type unification algorithm. unifier: Rc>>, + /// The fused JSON value selectors discovered while inferring `=`/`<>` over + /// encrypted JSON field accesses. Unification records *types* per node; this + /// records the one *relationship* the proxy needs — which operand supplies + /// the path for which value ([`crate::JsonValueSelectors`]). + json_value_selectors: RefCell>, + _ast: PhantomData<&'ast ()>, } @@ -65,10 +73,37 @@ impl<'ast> TypeInferencer<'ast> { table_resolver: table_resolver.into(), scope_tracker: scope.into(), unifier: unifier.into(), + json_value_selectors: RefCell::new(JsonValueSelectors::default()), _ast: PhantomData, } } + /// Takes the fused JSON value selectors accumulated during inference, + /// leaving the inferencer's set empty. + pub(crate) fn take_json_value_selectors(&self) -> JsonValueSelectors<'ast> { + std::mem::take(&mut self.json_value_selectors.borrow_mut()) + } + + pub(crate) fn record_json_value_selector_param( + &self, + param: Param, + source: JsonSelectorSource, + ) { + self.json_value_selectors + .borrow_mut() + .record_param(param, source); + } + + pub(crate) fn record_json_value_selector_literal( + &self, + node: &'ast sqltk::parser::ast::Value, + source: JsonSelectorSource, + ) { + self.json_value_selectors + .borrow_mut() + .record_literal(node, source); + } + pub(crate) fn get_node_type(&self, node: &'ast N) -> Arc { self.unifier.borrow_mut().get_node_type(node) } diff --git a/packages/eql-mapper/src/inference/unifier/eql_traits.rs b/packages/eql-mapper/src/inference/unifier/eql_traits.rs index 8ef5e4ce2..d02b69f70 100644 --- a/packages/eql-mapper/src/inference/unifier/eql_traits.rs +++ b/packages/eql-mapper/src/inference/unifier/eql_traits.rs @@ -327,6 +327,7 @@ impl EqlTerm { EqlTerm::JsonPath(_) => EqlTraits::none(), EqlTerm::Tokenized(_) => EqlTraits::none(), EqlTerm::JsonOrd(_) => EqlTraits::none(), + EqlTerm::JsonValueSelector(_) => EqlTraits::none(), } } } diff --git a/packages/eql-mapper/src/inference/unifier/types.rs b/packages/eql-mapper/src/inference/unifier/types.rs index 26eb1b0f8..646e6cb09 100644 --- a/packages/eql-mapper/src/inference/unifier/types.rs +++ b/packages/eql-mapper/src/inference/unifier/types.rs @@ -202,6 +202,19 @@ pub enum EqlTerm { /// compared via `eql_v3.ord_term`, so it is NOT a whole JSON document. #[display("EQL:JsonOrd({})", _0)] JsonOrd(EqlValue), + + /// The scalar value operand of an encrypted JSON field *equality* — the + /// non-JSON side of `col -> sel = value` (and the `->>` / + /// `jsonb_path_query_first` spellings). + /// + /// Exact JSON equality is selector containment, not a term comparison: the + /// needle is a single keyed MAC over `path ‖ canonical(value)` + /// (`QueryOp::SteVecValueSelector`), so this operand is **fused** from two + /// SQL operands — the path and the value — into one encrypted needle. The + /// mapper holds no encryption key, so it records only *where the path comes + /// from* ([`crate::JsonSelectorSource`]); the proxy composes and encrypts. + #[display("EQL:JsonValueSelector({})", _0)] + JsonValueSelector(EqlValue), } #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Display, Hash)] @@ -218,6 +231,8 @@ pub enum EqlTermVariant { Tokenized, #[display("EQL:JsonOrd")] JsonOrd, + #[display("EQL:JsonValueSelector")] + JsonValueSelector, } impl EqlTerm { @@ -234,7 +249,8 @@ impl EqlTerm { | EqlTerm::JsonAccessor(eql_value) | EqlTerm::JsonPath(eql_value) | EqlTerm::Tokenized(eql_value) - | EqlTerm::JsonOrd(eql_value) => eql_value, + | EqlTerm::JsonOrd(eql_value) + | EqlTerm::JsonValueSelector(eql_value) => eql_value, } } @@ -246,6 +262,7 @@ impl EqlTerm { EqlTerm::JsonPath(_) => EqlTermVariant::JsonPath, EqlTerm::Tokenized(_) => EqlTermVariant::Tokenized, EqlTerm::JsonOrd(_) => EqlTermVariant::JsonOrd, + EqlTerm::JsonValueSelector(_) => EqlTermVariant::JsonValueSelector, } } } diff --git a/packages/eql-mapper/src/inference/unifier/unify_types.rs b/packages/eql-mapper/src/inference/unifier/unify_types.rs index 57b9a0ae8..1a51f39cf 100644 --- a/packages/eql-mapper/src/inference/unifier/unify_types.rs +++ b/packages/eql-mapper/src/inference/unifier/unify_types.rs @@ -117,6 +117,10 @@ impl UnifyTypes for Unifier<'_> { Ok(EqlTerm::JsonOrd(lhs.clone()).into()) } + (EqlTerm::JsonValueSelector(lhs), EqlTerm::JsonValueSelector(rhs)) if lhs == rhs => { + Ok(EqlTerm::JsonValueSelector(lhs.clone()).into()) + } + (_, _) => Err(TypeError::Conflict(format!( "cannot unify EQL terms {lhs} and {rhs}" ))), diff --git a/packages/eql-mapper/src/json_value_selector.rs b/packages/eql-mapper/src/json_value_selector.rs new file mode 100644 index 000000000..df9ee6779 --- /dev/null +++ b/packages/eql-mapper/src/json_value_selector.rs @@ -0,0 +1,75 @@ +//! The N:1 fusion record for encrypted-JSON equality. +//! +//! `col -> sel = value` does not compare two encrypted terms. Exact JSON +//! equality in EQL v3 is *containment of a value selector*: a single keyed MAC +//! over the path and the canonicalised value together +//! (`QueryOp::SteVecValueSelector`, input `{"path": , "value": +//! }`). One needle, built from **two** SQL operands. +//! +//! The mapper cannot build it — it holds no encryption key. So the mapper does +//! the half it can: it types the value operand [`EqlTerm::JsonValueSelector`], +//! drops the path operand from the rewritten SQL, and records *where the path +//! came from* so the proxy can fuse the pair at encryption time. +//! +//! [`EqlTerm::JsonValueSelector`]: crate::EqlTerm::JsonValueSelector + +use std::collections::HashMap; + +use sqltk::parser::ast; +use sqltk::NodeKey; + +use crate::Param; + +/// Where the JSON path half of a fused value selector comes from. +/// +/// The two halves are independently a literal or a placeholder, so all four +/// combinations occur (`-> 'a' = '1'`, `-> $1 = $2`, `-> 'a' = $1`, …). A +/// literal path is fully known at type-check time and is carried inline; a +/// placeholder path is only known at Bind, so its param number is carried +/// instead. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum JsonSelectorSource { + /// A SQL literal path (`col -> 'name' = …`) — the selector text itself. + Literal(String), + + /// A placeholder path (`col -> $1 = …`) — the param it will arrive in. + Param(Param), +} + +/// The set of fused JSON value selectors in a statement: for each operand that +/// carries the *value* half, where its *path* half comes from. +/// +/// Keyed separately for the two protocols the proxy has to serve — params are +/// addressed by number (the extended protocol has no AST at Bind time), +/// literals by AST node. +#[derive(Debug, Default)] +pub struct JsonValueSelectors<'ast> { + by_param: HashMap, + by_literal: HashMap, JsonSelectorSource>, +} + +impl<'ast> JsonValueSelectors<'ast> { + pub(crate) fn record_param(&mut self, param: Param, source: JsonSelectorSource) { + self.by_param.insert(param, source); + } + + pub(crate) fn record_literal(&mut self, node: &'ast ast::Value, source: JsonSelectorSource) { + self.by_literal.insert(NodeKey::new(node), source); + } + + /// The path source for the value-selector operand bound to `param`, or + /// `None` if that param is not one. + pub fn for_param(&self, param: Param) -> Option<&JsonSelectorSource> { + self.by_param.get(¶m) + } + + /// The path source for the value-selector operand at literal `node`, or + /// `None` if that literal is not one. + pub fn for_literal(&self, node: &'ast ast::Value) -> Option<&JsonSelectorSource> { + self.by_literal.get(&NodeKey::new(node)) + } + + pub fn is_empty(&self) -> bool { + self.by_param.is_empty() && self.by_literal.is_empty() + } +} diff --git a/packages/eql-mapper/src/lib.rs b/packages/eql-mapper/src/lib.rs index e6030edf3..0f4d176ea 100644 --- a/packages/eql-mapper/src/lib.rs +++ b/packages/eql-mapper/src/lib.rs @@ -6,6 +6,7 @@ mod eql_mapper; mod importer; mod inference; mod iterator_ext; +mod json_value_selector; mod model; mod param; mod scope_tracker; @@ -17,6 +18,7 @@ mod test_helpers; pub use display_helpers::*; pub use eql_mapper::*; +pub use json_value_selector::*; pub use model::*; pub use param::*; pub use type_checked_statement::*; @@ -39,7 +41,7 @@ mod test { EqlTerm, EqlTrait, EqlTraits, EqlValue, InstantiateType, NativeValue, Projection, ProjectionColumn, Type, Value, }, - Param, Schema, TableColumn, TableResolver, + JsonSelectorSource, Param, Schema, TableColumn, TableResolver, }; use eql_mapper_macros::concrete_ty; use pretty_assertions::assert_eq; @@ -2349,6 +2351,135 @@ mod test { } } + /// A schema with one encrypted JSON column, for the value-selector tests. + /// + /// The domain is spelled out: value-selector fusion keys off the column's + /// *token* type being `Json`, and the macro's default synthesises a `text` + /// token regardless of the `JsonLike` capability. + fn json_eq_schema() -> Arc { + resolver(schema! { + tables: { + patients: { + id, + notes (EQL("eql_v3_json_search"): JsonLike + Contain), + } + } + }) + } + + /// `col -> $1 = $2` fuses both operands into ONE containment needle: the + /// field access is discarded, the value operand becomes the needle, and the + /// selector placeholder is left declared-but-unreferenced (which keeps param + /// numbering identical between client and server). + #[test] + fn json_field_eq_params_rewrites_to_containment() { + let statement = parse("SELECT id FROM patients WHERE notes -> $1 = $2"); + + let typed = type_check(json_eq_schema(), &statement).unwrap(); + + assert_eq!( + typed.transform(HashMap::new()).unwrap().to_string(), + "SELECT id FROM patients WHERE eql_v3.jsonb_contains(notes, $2::JSONB::eql_v3.query_json)" + ); + + // The value operand ($2) carries the needle; its path comes from $1. + assert_eq!( + typed.json_value_selectors.for_param(Param(2)), + Some(&JsonSelectorSource::Param(Param(1))) + ); + assert_eq!(typed.json_value_selectors.for_param(Param(1)), None); + } + + /// The `->>` and `jsonb_path_query_first` spellings are the same access and + /// rewrite identically. + #[test] + fn json_field_eq_alternate_spellings_rewrite_to_containment() { + for access in ["notes ->> $1", "jsonb_path_query_first(notes, $1)"] { + let statement = parse(&format!("SELECT id FROM patients WHERE {access} = $2")); + let typed = type_check(json_eq_schema(), &statement).unwrap(); + + assert_eq!( + typed.transform(HashMap::new()).unwrap().to_string(), + "SELECT id FROM patients WHERE eql_v3.jsonb_contains(notes, $2::JSONB::eql_v3.query_json)", + "unexpected rewrite for `{access}`" + ); + } + } + + /// Both halves as literals: the selector text is captured inline at + /// type-check time, and the selector literal vanishes from the output. + #[test] + fn json_field_eq_literals_rewrites_to_containment() { + let statement = + parse("SELECT id FROM patients WHERE notes -> 'medications' = '\"aspirin\"'"); + + let typed = type_check(json_eq_schema(), &statement).unwrap(); + + assert_eq!( + typed + .json_value_selectors + .for_literal(&ast::Value::SingleQuotedString("\"aspirin\"".to_owned())), + None, + "for_literal is keyed by node identity, not by value" + ); + + // Both literals are encrypted operands; only the value one survives the + // rewrite, so the selector's replacement is simply never placed. + let encrypted = HashMap::from_iter([ + ( + test_helpers::get_node_key_of_json_selector( + &statement, + &ast::Value::SingleQuotedString("medications".to_owned()), + ), + ast::Value::SingleQuotedString("".to_owned()), + ), + ( + test_helpers::get_node_key_of_json_selector( + &statement, + &ast::Value::SingleQuotedString("\"aspirin\"".to_owned()), + ), + ast::Value::SingleQuotedString("".to_owned()), + ), + ]); + + assert_eq!( + typed.transform(encrypted).unwrap().to_string(), + "SELECT id FROM patients WHERE eql_v3.jsonb_contains(notes, ''::JSONB::eql_v3.query_json)" + ); + } + + /// `<>` is containment negated. + #[test] + fn json_field_not_eq_rewrites_to_negated_containment() { + let statement = parse("SELECT id FROM patients WHERE notes -> $1 <> $2"); + + let typed = type_check(json_eq_schema(), &statement).unwrap(); + + assert_eq!( + typed.transform(HashMap::new()).unwrap().to_string(), + "SELECT id FROM patients WHERE NOT (eql_v3.jsonb_contains(notes, $2::JSONB::eql_v3.query_json))" + ); + } + + /// Equality on the whole encrypted JSON column is document equality, NOT a + /// field access — it must keep its ordinary term rewrite. Guards against the + /// value-selector fusion swallowing every `=` on a JSON column. + #[test] + fn json_column_eq_is_not_value_selector_containment() { + let statement = parse("SELECT id FROM patients WHERE notes = $1"); + + let typed = type_check(json_eq_schema(), &statement).unwrap(); + + assert_eq!(typed.json_value_selectors.for_param(Param(1)), None); + assert!(typed.json_value_selectors.is_empty()); + + let sql = typed.transform(HashMap::new()).unwrap().to_string(); + assert!( + !sql.contains("jsonb_contains"), + "whole-column equality must not become containment, got: {sql}" + ); + } + #[test] fn jsonb_path_query_param_to_eql() { // init_tracing(); diff --git a/packages/eql-mapper/src/transformation_rules/cast_literals_as_encrypted.rs b/packages/eql-mapper/src/transformation_rules/cast_literals_as_encrypted.rs index ea51b5814..53e68032f 100644 --- a/packages/eql-mapper/src/transformation_rules/cast_literals_as_encrypted.rs +++ b/packages/eql-mapper/src/transformation_rules/cast_literals_as_encrypted.rs @@ -7,7 +7,7 @@ use sqltk::{NodeKey, NodePath, Visitable}; use crate::unifier::{EqlTerm, Type, Value as UnifierValue}; use crate::EqlMapperError; -use super::helpers::{cast_to_v3_domain, json_ord_cast_target, v3_cast_target}; +use super::helpers::{cast_to_v3_domain, json_query_operand_cast_target, v3_cast_target}; use super::TransformationRule; #[derive(Debug)] @@ -67,8 +67,8 @@ impl<'ast> TransformationRule<'ast> for CastLiteralsAsEncrypted<'ast> { }) } _ => { - let (schema, domain) = - json_ord_cast_target(eql_term).unwrap_or_else(|| { + let (schema, domain) = json_query_operand_cast_target(eql_term) + .unwrap_or_else(|| { let identity = eql_term.eql_value().domain_identity().clone(); v3_cast_target(node_path, &identity) }); diff --git a/packages/eql-mapper/src/transformation_rules/cast_params_as_encrypted.rs b/packages/eql-mapper/src/transformation_rules/cast_params_as_encrypted.rs index f7259f582..56278b690 100644 --- a/packages/eql-mapper/src/transformation_rules/cast_params_as_encrypted.rs +++ b/packages/eql-mapper/src/transformation_rules/cast_params_as_encrypted.rs @@ -1,4 +1,4 @@ -use super::helpers::{cast_to_v3_domain, json_ord_cast_target, v3_cast_target}; +use super::helpers::{cast_to_v3_domain, json_query_operand_cast_target, v3_cast_target}; use super::TransformationRule; use crate::unifier::{EqlTerm, Type, Value as UnifierValue}; use crate::EqlMapperError; @@ -49,7 +49,7 @@ impl<'ast> TransformationRule<'ast> for CastParamsAsEncrypted<'ast> { return Ok(false); } - let (schema, domain) = json_ord_cast_target(eql_term).unwrap_or_else(|| { + let (schema, domain) = json_query_operand_cast_target(eql_term).unwrap_or_else(|| { let identity = eql_term.eql_value().domain_identity().clone(); v3_cast_target(node_path, &identity) }); diff --git a/packages/eql-mapper/src/transformation_rules/helpers.rs b/packages/eql-mapper/src/transformation_rules/helpers.rs index 5846f5cf7..b8b8cf711 100644 --- a/packages/eql-mapper/src/transformation_rules/helpers.rs +++ b/packages/eql-mapper/src/transformation_rules/helpers.rs @@ -19,9 +19,29 @@ use crate::unifier::{DomainIdentity, EqlTerm}; /// text, dates, etc. — which is what makes JSON range work in the extended /// protocol, where the operand's scalar type is unknown at rewrite time. /// Returns `None` for any other term. -pub(crate) fn json_ord_cast_target(eql_term: &EqlTerm) -> Option<(String, String)> { - matches!(eql_term, EqlTerm::JsonOrd(_)) - .then(|| ("eql_v3".to_string(), "query_integer_ord".to_string())) +/// The v3 cast target `(schema, domain)` for an encrypted-JSON *query operand* +/// whose domain is fixed by the operand's role rather than by the column's +/// domain identity. Returns `None` for any other term. +/// +/// - [`EqlTerm::JsonOrd`] — the shape-only scalar ord twin +/// `eql_v3.query_integer_ord`, regardless of the JSON leaf's scalar type. +/// `eql_v3.ord_term` is type-agnostic (it extracts the `op` bytes as +/// `ope_cllw` and compares them bytewise) and the twin's domain CHECK is +/// shape-only (`{v,i,op}`, no `c`), so one twin serves numbers, text, dates. +/// That is what makes JSON range work in the extended protocol, where the +/// operand's scalar type is unknown at rewrite time. +/// - [`EqlTerm::JsonValueSelector`] — `eql_v3.query_json`, the containment +/// needle domain. The fused value selector is a one-entry, term-less +/// containment payload (`{sv: [{s}]}`), which is what +/// `eql_v3.jsonb_contains` matches against. +pub(crate) fn json_query_operand_cast_target(eql_term: &EqlTerm) -> Option<(String, String)> { + let domain = match eql_term { + EqlTerm::JsonOrd(_) => "query_integer_ord", + EqlTerm::JsonValueSelector(_) => "query_json", + _ => return None, + }; + + Some(("eql_v3".to_string(), domain.to_string())) } /// The scalar comparison operators the v3 term-function rewrite handles. diff --git a/packages/eql-mapper/src/transformation_rules/mod.rs b/packages/eql-mapper/src/transformation_rules/mod.rs index d48cab5bf..4ac1f4007 100644 --- a/packages/eql-mapper/src/transformation_rules/mod.rs +++ b/packages/eql-mapper/src/transformation_rules/mod.rs @@ -18,6 +18,7 @@ mod preserve_effective_aliases; mod rewrite_containment_ops; mod rewrite_eql_comparison_ops; mod rewrite_eql_match_ops; +mod rewrite_json_value_selector_eq; mod rewrite_standard_sql_fns_on_eql_types; use std::marker::PhantomData; @@ -29,6 +30,7 @@ pub(crate) use preserve_effective_aliases::*; pub(crate) use rewrite_containment_ops::*; pub(crate) use rewrite_eql_comparison_ops::*; pub(crate) use rewrite_eql_match_ops::*; +pub(crate) use rewrite_json_value_selector_eq::*; pub(crate) use rewrite_standard_sql_fns_on_eql_types::*; use crate::EqlMapperError; diff --git a/packages/eql-mapper/src/transformation_rules/rewrite_eql_comparison_ops.rs b/packages/eql-mapper/src/transformation_rules/rewrite_eql_comparison_ops.rs index 042b04ecb..258e527f1 100644 --- a/packages/eql-mapper/src/transformation_rules/rewrite_eql_comparison_ops.rs +++ b/packages/eql-mapper/src/transformation_rules/rewrite_eql_comparison_ops.rs @@ -7,7 +7,7 @@ use sqltk::parser::ast::{BinaryOperator, Expr, ValueWithSpan}; use sqltk::parser::tokenizer::Span; use sqltk::{NodeKey, NodePath, Visitable}; -use crate::unifier::{DomainIdentity, Type, Value}; +use crate::unifier::{DomainIdentity, EqlTerm, Type, Value}; use crate::EqlMapperError; use super::helpers::{eql_v3_term_call, is_comparison_op}; @@ -49,6 +49,19 @@ impl<'ast> RewriteEqlComparisonOps<'ast> { } } + /// Encrypted JSON field equality is value-selector containment, rewritten by + /// [`super::RewriteJsonValueSelectorEq`], not a term comparison. `eq_term` + /// has no unique overload for a JSON query operand, so wrapping one here + /// would produce SQL PostgreSQL rejects. + fn is_json_value_selector_eq(&self, left: &'ast Expr, right: &'ast Expr) -> bool { + [left, right].into_iter().any(|expr| { + matches!( + self.node_types.get(&NodeKey::new(expr)), + Some(Type::Value(Value::Eql(EqlTerm::JsonValueSelector(_)))) + ) + }) + } + /// The term function for `op` on a column with `identity`, or `None` if the /// domain provides no term for that operator. fn term_fn_for(op: &BinaryOperator, identity: &DomainIdentity) -> Option<&'static str> { @@ -79,7 +92,7 @@ impl<'ast> TransformationRule<'ast> for RewriteEqlComparisonOps<'ast> { let Some((Expr::BinaryOp { left, op, right },)) = node_path.last_1_as::() else { return Ok(false); }; - if !is_comparison_op(op) { + if !is_comparison_op(op) || self.is_json_value_selector_eq(left, right) { return Ok(false); } let Some(identity) = self @@ -113,7 +126,7 @@ impl<'ast> TransformationRule<'ast> for RewriteEqlComparisonOps<'ast> { fn would_edit(&mut self, node_path: &NodePath<'ast>, _target_node: &N) -> bool { if let Some((Expr::BinaryOp { left, op, right },)) = node_path.last_1_as::() { - if is_comparison_op(op) { + if is_comparison_op(op) && !self.is_json_value_selector_eq(left, right) { return self.eql_identity_of(left).is_some() || self.eql_identity_of(right).is_some(); } diff --git a/packages/eql-mapper/src/transformation_rules/rewrite_json_value_selector_eq.rs b/packages/eql-mapper/src/transformation_rules/rewrite_json_value_selector_eq.rs new file mode 100644 index 000000000..8af09435f --- /dev/null +++ b/packages/eql-mapper/src/transformation_rules/rewrite_json_value_selector_eq.rs @@ -0,0 +1,195 @@ +use std::collections::HashMap; +use std::mem; +use std::sync::Arc; + +use sqltk::parser::ast::Value as SqltkValue; +use sqltk::parser::ast::{ + BinaryOperator, Expr, Function, FunctionArg, FunctionArgExpr, FunctionArgumentList, + FunctionArguments, Ident, ObjectName, ObjectNamePart, ValueWithSpan, +}; +use sqltk::parser::tokenizer::Span; +use sqltk::{NodeKey, NodePath, Visitable}; + +use crate::unifier::{EqlTerm, Type, Value}; +use crate::EqlMapperError; + +use super::TransformationRule; + +/// Rewrites equality on an encrypted JSON **field** into value-selector +/// containment: +/// +/// - `col -> sel = value` → `eql_v3.jsonb_contains(col, )` +/// - `col ->> sel = value` → same +/// - `jsonb_path_query_first(col, sel) = value` → same +/// - `<>` negates: `NOT eql_v3.jsonb_contains(col, )` +/// +/// where `` is the value operand, already cast to `eql_v3.query_json` +/// by the cast rules. Exact JSON equality in EQL v3 is selector containment: the +/// needle is one keyed MAC over the path and the canonicalised value together +/// (`QueryOp::SteVecValueSelector`), which the proxy composes from the two SQL +/// operands. So this rule **discards** the field access: `col` is lifted out and +/// the selector operand disappears from the statement. +/// +/// A discarded selector *placeholder* stays declared in Parse but unreferenced +/// in the SQL, which PostgreSQL permits as long as its type is known. That keeps +/// input and output param numbering identical, so Bind stays positional. +/// +/// The generic `eq_term` wrap ([`super::RewriteEqlComparisonOps`]) must not also +/// fire here — `eql_v3.eq_term` has no unique overload for a JSON query operand, +/// and containment is not a term comparison. That rule skips any comparison +/// whose operand is [`EqlTerm::JsonValueSelector`]. +#[derive(Debug)] +pub struct RewriteJsonValueSelectorEq<'ast> { + node_types: Arc, Type>>, +} + +impl<'ast> RewriteJsonValueSelectorEq<'ast> { + pub fn new(node_types: Arc, Type>>) -> Self { + Self { node_types } + } + + /// Whether `expr` is the fused value operand of a JSON field equality. + fn is_value_selector(&self, expr: &'ast Expr) -> bool { + matches!( + self.node_types.get(&NodeKey::new(expr)), + Some(Type::Value(Value::Eql(EqlTerm::JsonValueSelector(_)))) + ) + } + + /// The container expression of a JSON field access — the `col` of + /// `col -> sel` or of `jsonb_path_query_first(col, sel)`. + /// + /// Read from the ORIGINAL AST (via `node_path`), because by the time this + /// rule runs the field access has already been rewritten to + /// `eql_v3."->"(col, sel)` by [`super::RewriteContainmentOps`]. + fn container_of(expr: &Expr) -> Option<&Expr> { + match expr { + Expr::BinaryOp { + left, + op: BinaryOperator::Arrow | BinaryOperator::LongArrow, + .. + } => Some(&**left), + + Expr::Function(function) => match &function.args { + FunctionArguments::List(list) => match list.args.as_slice() { + [FunctionArg::Unnamed(FunctionArgExpr::Expr(container)), _] => Some(container), + _ => None, + }, + _ => None, + }, + + _ => None, + } + } + + /// Splits a comparison into `(container, value operand is on the right)`, or + /// `None` if this is not a JSON field equality. + fn match_comparison( + &self, + left: &'ast Expr, + op: &BinaryOperator, + right: &'ast Expr, + ) -> Option<(&'ast Expr, bool)> { + if !matches!(op, BinaryOperator::Eq | BinaryOperator::NotEq) { + return None; + } + + if self.is_value_selector(right) { + Self::container_of(left).map(|container| (container, true)) + } else if self.is_value_selector(left) { + Self::container_of(right).map(|container| (container, false)) + } else { + None + } + } + + fn jsonb_contains(container: Expr, needle: Expr) -> Expr { + Expr::Function(Function { + name: ObjectName(vec![ + ObjectNamePart::Identifier(Ident::new("eql_v3")), + ObjectNamePart::Identifier(Ident::new("jsonb_contains")), + ]), + uses_odbc_syntax: false, + args: FunctionArguments::List(FunctionArgumentList { + args: vec![ + FunctionArg::Unnamed(FunctionArgExpr::Expr(container)), + FunctionArg::Unnamed(FunctionArgExpr::Expr(needle)), + ], + duplicate_treatment: None, + clauses: vec![], + }), + parameters: FunctionArguments::None, + filter: None, + null_treatment: None, + over: None, + within_group: vec![], + }) + } +} + +impl<'ast> TransformationRule<'ast> for RewriteJsonValueSelectorEq<'ast> { + fn apply( + &mut self, + node_path: &NodePath<'ast>, + target_node: &mut N, + ) -> Result { + // Match against the ORIGINAL nodes: `node_types` is keyed by them, and + // `target_node`'s children have already been rebuilt by earlier rules. + let Some((Expr::BinaryOp { left, op, right },)) = node_path.last_1_as::() else { + return Ok(false); + }; + + let Some((container, value_on_right)) = self.match_comparison(left, op, right) else { + return Ok(false); + }; + + let negated = matches!(op, BinaryOperator::NotEq); + + let Some(expr) = target_node.downcast_mut::() else { + return Ok(false); + }; + let Expr::BinaryOp { + left: target_left, + right: target_right, + .. + } = expr + else { + return Ok(false); + }; + + // Move (not clone) the transformed value operand so its NodeKey identity + // survives; the container comes from the original AST, where it is still + // the bare column reference the containment call needs. + let dummy = Expr::Value(ValueWithSpan { + value: SqltkValue::Null, + span: Span::empty(), + }); + let needle = if value_on_right { + mem::replace(&mut **target_right, dummy) + } else { + mem::replace(&mut **target_left, dummy) + }; + + let contains = Self::jsonb_contains(container.clone(), needle); + + *expr = if negated { + Expr::UnaryOp { + op: sqltk::parser::ast::UnaryOperator::Not, + expr: Box::new(Expr::Nested(Box::new(contains))), + } + } else { + contains + }; + + Ok(true) + } + + fn would_edit(&mut self, node_path: &NodePath<'ast>, _target_node: &N) -> bool { + match node_path.last_1_as::() { + Some((Expr::BinaryOp { left, op, right },)) => { + self.match_comparison(left, op, right).is_some() + } + _ => false, + } + } +} diff --git a/packages/eql-mapper/src/type_checked_statement.rs b/packages/eql-mapper/src/type_checked_statement.rs index e46cc10b4..d462f4675 100644 --- a/packages/eql-mapper/src/type_checked_statement.rs +++ b/packages/eql-mapper/src/type_checked_statement.rs @@ -6,9 +6,9 @@ use sqltk::{AsNodeKey, NodeKey, Transformable}; use crate::unifier::EqlTerm; use crate::{ CastLiteralsAsEncrypted, CastParamsAsEncrypted, DryRunnable, EqlMapperError, - FailOnPlaceholderChange, Param, PreserveEffectiveAliases, RewriteContainmentOps, - RewriteEqlComparisonOps, RewriteEqlMatchOps, RewriteStandardSqlFnsOnEqlTypes, - TransformationRule, + FailOnPlaceholderChange, JsonValueSelectors, Param, PreserveEffectiveAliases, + RewriteContainmentOps, RewriteEqlComparisonOps, RewriteEqlMatchOps, RewriteJsonValueSelectorEq, + RewriteStandardSqlFnsOnEqlTypes, TransformationRule, }; use crate::unifier::{Projection, Type, Value}; @@ -28,6 +28,17 @@ pub struct TypeCheckedStatement<'ast> { /// The type ([`EqlTerm`]) and reference to an [`ast::Value`] nodes of all EQL literals from the SQL statement. pub literals: Vec<(EqlTerm, &'ast ast::Value)>, + /// The fused JSON value selectors: for each operand typed + /// [`EqlTerm::JsonValueSelector`], where the path half of its needle comes + /// from. + /// + /// This is the one place the mapper's output is **not** 1:1 with the input + /// SQL. `col -> sel = value` consumes two operands and emits one encrypted + /// needle; the path operand is dropped from the rewritten statement (its + /// placeholder stays declared but unreferenced, so param numbering is + /// untouched). The proxy consults this to compose the needle. + pub json_value_selectors: JsonValueSelectors<'ast>, + /// A [`HashMap`] of AST node (using [`NodeKey`] as the key) to [`Type`]. The map contains a `Type` for every node /// in the AST with the node type is one of: [`Statement`], [`Query`], [`Insert`], [`Delete`], [`Expr`], /// [`SetExpr`], [`Select`], [`SelectItem`], [`Vec`], [`Function`], [`Values`], [`Value`]. @@ -52,6 +63,7 @@ impl<'ast> TypeCheckedStatement<'ast> { projection: Projection, params: Vec<(Param, Value)>, literals: Vec<(EqlTerm, &'ast ast::Value)>, + json_value_selectors: JsonValueSelectors<'ast>, node_types: Arc, Type>>, ) -> Self { Self { @@ -59,6 +71,7 @@ impl<'ast> TypeCheckedStatement<'ast> { projection, params, literals, + json_value_selectors, node_types, } } @@ -154,6 +167,7 @@ impl<'ast> TypeCheckedStatement<'ast> { DryRunnable::new(( RewriteStandardSqlFnsOnEqlTypes::new(Arc::clone(&self.node_types)), RewriteContainmentOps::new(Arc::clone(&self.node_types)), + RewriteJsonValueSelectorEq::new(Arc::clone(&self.node_types)), RewriteEqlComparisonOps::new(Arc::clone(&self.node_types)), RewriteEqlMatchOps::new(Arc::clone(&self.node_types)), PreserveEffectiveAliases, From 474af85e9ba442316f37d34b968b6e77f2f6f3a9 Mon Sep 17 00:00:00 2001 From: James Sadler Date: Sun, 26 Jul 2026 01:10:10 +1000 Subject: [PATCH 02/10] feat(proxy): compose and encrypt the JSON equality value selector MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The proxy half of encrypted JSON field equality. The mapper types the value operand `JsonValueSelector` and records where its path comes from; this fuses the two SQL operands into one `{"path", "value"}` plaintext and encrypts it with `QueryOp::SteVecValueSelector`, producing the one-entry containment needle `eql_v3.jsonb_contains` matches. Both protocols: - Simple: `literals_to_plaintext` composes the needle for a value-selector literal from its own JSON scalar and the selector literal's text. The discarded selector literal is still encrypted but never placed, since the rewrite removed its node. - Extended: `Bind::to_plaintext` gains a second pass. A fused param is skipped by the positional pass — its own bytes are only half a needle and would not decode as a standalone operand for the column — then composed from the path param (or literal) and its own value. - `Parse::declare_unreferenced_param_types` declares TEXT for the selector placeholder the rewrite dropped. Clients that let the server infer types send no types at all, and PostgreSQL will not prepare a statement with an unreferenced param of unknown type. Still-referenced slots are left at OID 0 ("infer"), so this never overrides a type the query determines. Param counts are unchanged end to end, so Bind, ParameterDescription and the encrypt pipeline all stay positional. Also folds the four copies of the `$.`-rooting of a JSON selector into one `json_selector_path` helper — the fused path needs the same normalisation the accessor and path-query operands already did. Fixes `select_jsonb_where_{string,numeric}_eq` (both protocols, both the `->` and `jsonb_path_query_first` spellings). jsonb suite 70 -> 72 passing. Verified zero regressions by running the full integration suite against the base binary and diffing failure sets: the only differences are these two tests plus known-flaky ORE ordering tests that fail identically on base. The 7 `jsonb_term_filter` tests remain red for an unrelated reason: `encrypted_jsonb_filtered` is declared `eql_v3_json_search`, the same domain as the unfiltered column, and `from_domain.rs` hardcodes `term_filters: Vec::new()`. EQL v3 has no domain or config that encodes a term filter, so the proxy cannot infer one. Equality itself works on that column — it matches case-sensitively. Case-insensitive JSON needs an EQL-side way to declare filters; that is a separate feature. Stable-Commit-Id: q-5afqdbdnjhd299x8777mwdm4fj --- .../src/postgresql/context/mod.rs | 7 +- .../src/postgresql/context/statement.rs | 38 +++++ .../src/postgresql/data/from_sql.rs | 140 +++++++++++++++--- .../src/postgresql/data/mod.rs | 1 + .../src/postgresql/frontend.rs | 108 +++++++++++--- .../src/postgresql/messages/bind.rs | 86 ++++++++++- .../src/postgresql/messages/parse.rs | 36 +++++ .../src/proxy/zerokms/zerokms.rs | 17 +++ 8 files changed, 391 insertions(+), 42 deletions(-) diff --git a/packages/cipherstash-proxy/src/postgresql/context/mod.rs b/packages/cipherstash-proxy/src/postgresql/context/mod.rs index f9cd0ebe0..50411dc47 100644 --- a/packages/cipherstash-proxy/src/postgresql/context/mod.rs +++ b/packages/cipherstash-proxy/src/postgresql/context/mod.rs @@ -3,7 +3,11 @@ pub mod phase_timing; pub mod portal; pub mod statement; pub mod statement_metadata; -pub use self::{phase_timing::PhaseTiming, portal::Portal, statement::Statement}; +pub use self::{ + phase_timing::PhaseTiming, + portal::Portal, + statement::{JsonSelectorPath, Statement}, +}; use super::{ column_mapper::ColumnMapper, messages::{describe::Describe, Name, Target}, @@ -1116,6 +1120,7 @@ mod tests { projection_columns: vec![], literal_columns: vec![], postgres_param_types: vec![], + json_value_selectors: std::collections::HashMap::new(), } } diff --git a/packages/cipherstash-proxy/src/postgresql/context/statement.rs b/packages/cipherstash-proxy/src/postgresql/context/statement.rs index 68170d823..6bd90d6df 100644 --- a/packages/cipherstash-proxy/src/postgresql/context/statement.rs +++ b/packages/cipherstash-proxy/src/postgresql/context/statement.rs @@ -1,4 +1,19 @@ use super::Column; +use std::collections::HashMap; + +/// Where the path half of a fused JSON value selector comes from, resolved to +/// this statement's bind params. +/// +/// The proxy's copy of [`eql_mapper::JsonSelectorSource`], with param numbers +/// converted to 0-based bind indexes. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum JsonSelectorPath { + /// A literal path in the SQL, known at Parse time. + Literal(String), + + /// A placeholder path, arriving in this (0-based) bind param. + Param(usize), +} /// /// Type Analysed parameters and projection @@ -9,6 +24,13 @@ pub struct Statement { pub projection_columns: Vec>, pub literal_columns: Vec>, pub postgres_param_types: Vec, + + /// Params that carry a fused JSON value-selector needle (`col -> sel = + /// $n`), keyed by 0-based bind index, mapped to where their path comes from. + /// + /// Bind consults this to compose `{"path", "value"}` from two params before + /// encrypting. Empty for every statement without encrypted JSON equality. + pub json_value_selectors: HashMap, } impl Statement { @@ -17,15 +39,31 @@ impl Statement { projection_columns: Vec>, literal_columns: Vec>, postgres_param_types: Vec, + json_value_selectors: HashMap, ) -> Statement { Statement { param_columns, projection_columns, literal_columns, postgres_param_types, + json_value_selectors, } } + /// The 0-based bind indexes of params that supply only a value-selector + /// path. The rewrite folds these into the needle and drops them from the + /// SQL, so PostgreSQL never references them — but it still needs their type + /// declared in Parse to prepare the statement. + pub fn unreferenced_param_indexes(&self) -> Vec { + self.json_value_selectors + .values() + .filter_map(|path| match path { + JsonSelectorPath::Param(idx) => Some(*idx), + JsonSelectorPath::Literal(_) => None, + }) + .collect() + } + pub fn has_literals(&self) -> bool { !self.literal_columns.is_empty() } diff --git a/packages/cipherstash-proxy/src/postgresql/data/from_sql.rs b/packages/cipherstash-proxy/src/postgresql/data/from_sql.rs index 1b3f757ae..0bb09fa8b 100644 --- a/packages/cipherstash-proxy/src/postgresql/data/from_sql.rs +++ b/packages/cipherstash-proxy/src/postgresql/data/from_sql.rs @@ -106,6 +106,119 @@ pub fn literal_from_sql( Ok(pt) } +/// Normalises a JSON field selector to an eJSONPath rooted at `$`. +/// +/// `->`/`->>` take a bare field name (`name`), `jsonb_path_query*` takes a path +/// (`nested.title`, or already-rooted `$.nested.title`). The client's +/// `Selector::parse` only accepts the rooted form. +pub fn json_selector_path(val: &str) -> String { + if val.starts_with("$.") { + val.to_string() + } else { + format!("$.{val}") + } +} + +/// Builds the composition input for a fused JSON value selector: +/// `{"path": , "value": }`. +/// +/// This is the one place two SQL operands become one encrypted operand. +/// `QueryOp::SteVecValueSelector` MACs the path and the canonicalised value +/// together into a single selector; its presence in the stored `sv` is the +/// equality match. The client applies the column's term filters (e.g. downcase) +/// to `value` as part of that, so case-insensitive columns work unchanged here. +/// +/// `value` must be a scalar. A single value selector is only injective for +/// scalars — a container MACs just its structural tag, so every object at a path +/// would collapse to one selector. The client rejects those; rejecting here too +/// gives a message naming the query shape rather than the encryption internals. +pub fn json_value_selector_plaintext( + path: &str, + value: serde_json::Value, +) -> Result { + if value.is_object() || value.is_array() { + debug!( + target: ENCODING, + msg = "Encrypted JSON equality requires a scalar value", + ?path, + ?value + ); + return Err(MappingError::CouldNotParseParameter); + } + + Ok(Plaintext::new(serde_json::json!({ + "path": json_selector_path(path), + "value": value, + }))) +} + +/// The JSON value a literal carries, for fusing into a value selector. +/// +/// The value half of `col -> sel = value` is written as a quoted JSON scalar +/// (`= '"B"'`, `= '3'`) or as a bare SQL number (`= 3`). A quoted string that is +/// not valid JSON is taken as the string itself, so `= 'B'` behaves like +/// `= '"B"'`. +pub fn literal_json_value(literal: &Value) -> Result, MappingError> { + let value = match literal { + Value::Null => None, + + Value::Number(d, _) => Some( + serde_json::from_str::(&d.to_string()) + .map_err(|_| MappingError::CouldNotParseParameter)?, + ), + + Value::Boolean(b) => Some(serde_json::Value::Bool(*b)), + + Value::SingleQuotedString(s) + | Value::DoubleQuotedString(s) + | Value::TripleSingleQuotedString(s) + | Value::TripleDoubleQuotedString(s) + | Value::EscapedStringLiteral(s) + | Value::UnicodeStringLiteral(s) + | Value::NationalStringLiteral(s) => Some( + serde_json::from_str::(s) + .unwrap_or_else(|_| serde_json::Value::String(s.to_owned())), + ), + + Value::DollarQuotedString(s) => Some( + serde_json::from_str::(&s.value) + .unwrap_or_else(|_| serde_json::Value::String(s.value.to_owned())), + ), + + _ => return Err(MappingError::CouldNotParseParameter), + }; + + Ok(value) +} + +/// The JSON value a bind param carries, for fusing into a value selector. +/// +/// Mirrors [`literal_json_value`] for the extended protocol: a jsonb param +/// arrives either as its text rendering (`4`, `"C"`) or as binary jsonb. A text +/// payload that is not valid JSON is taken as the string itself. +pub fn bind_param_json_value( + param: &BindParam, + postgres_type: &Type, +) -> Result, MappingError> { + if param.is_null() { + return Ok(None); + } + + let value = match param.format_code { + FormatCode::Text => { + let text = param.to_string(); + serde_json::from_str::(&text) + .unwrap_or(serde_json::Value::String(text)) + } + FormatCode::Binary => { + parse_bytes_from_sql::(¶m.bytes, postgres_type) + .map_err(|_| MappingError::CouldNotParseParameter)? + } + }; + + Ok(Some(value)) +} + /// A JSON ordering operand (`EqlTerm::JsonOrd`) arriving as a jsonb param /// carries a single scalar. Encode a number as a float (matching the stored /// JSON number leaf's SteVec `op` encoding) and a string as text — the only @@ -197,12 +310,7 @@ fn text_from_sql( // If JSONB, JSONPATH values are treated as strings (EqlTermVariant::JsonPath | EqlTermVariant::JsonAccessor, ColumnType::Json) => { - let val = if val.starts_with("$.") { - val.to_string() - } else { - format!("$.{val}") - }; - Ok(Plaintext::new(val)) + Ok(Plaintext::new(json_selector_path(val))) } (EqlTermVariant::Full | EqlTermVariant::Partial, ColumnType::Json) => { serde_json::from_str::(val) @@ -297,24 +405,12 @@ fn binary_from_sql( // If JSONB, JSONPATH values are treated as strings (EqlTermVariant::JsonPath, ColumnType::Json, &Type::JSONPATH) => { - parse_bytes_from_sql::(bytes, pg_type).map(|val| { - let val = if val.starts_with("$.") { - val - } else { - format!("$.{val}") - }; - Plaintext::new(val) - }) + parse_bytes_from_sql::(bytes, pg_type) + .map(|val| Plaintext::new(json_selector_path(&val))) } (EqlTermVariant::JsonAccessor, ColumnType::Json, &Type::TEXT | &Type::VARCHAR) => { - parse_bytes_from_sql::(bytes, pg_type).map(|val| { - let val = if val.starts_with("$.") { - val - } else { - format!("$.{val}") - }; - Plaintext::new(val) - }) + parse_bytes_from_sql::(bytes, pg_type) + .map(|val| Plaintext::new(json_selector_path(&val))) } // A JSON ordering operand (`col -> sel > $2`) arrives as a jsonb scalar; // encode it as the scalar shape SteVecTerm accepts (number → float, diff --git a/packages/cipherstash-proxy/src/postgresql/data/mod.rs b/packages/cipherstash-proxy/src/postgresql/data/mod.rs index 6dc03f075..63c618c75 100644 --- a/packages/cipherstash-proxy/src/postgresql/data/mod.rs +++ b/packages/cipherstash-proxy/src/postgresql/data/mod.rs @@ -9,6 +9,7 @@ use rust_decimal::{prelude::FromPrimitive, Decimal}; use tracing::{debug, warn}; pub use from_sql::literal_from_sql; +pub use from_sql::{bind_param_json_value, json_value_selector_plaintext, literal_json_value}; pub use to_sql::to_sql; /// /// Fun fact: some clients can specify a parameter type with a parse message diff --git a/packages/cipherstash-proxy/src/postgresql/frontend.rs b/packages/cipherstash-proxy/src/postgresql/frontend.rs index b3e7822d2..e6d328516 100644 --- a/packages/cipherstash-proxy/src/postgresql/frontend.rs +++ b/packages/cipherstash-proxy/src/postgresql/frontend.rs @@ -14,8 +14,11 @@ use crate::error::{EncryptError, Error, MappingError}; use crate::log::{MAPPER, PROTOCOL}; use crate::postgresql::context::column::Column; use crate::postgresql::context::statement_metadata::{ProtocolType, StatementType}; +use crate::postgresql::context::JsonSelectorPath; use crate::postgresql::context::Portal; -use crate::postgresql::data::literal_from_sql; +use crate::postgresql::data::{ + json_value_selector_plaintext, literal_from_sql, literal_json_value, +}; use crate::postgresql::messages::close::Close; use crate::postgresql::messages::ready_for_query::ReadyForQuery; use crate::postgresql::messages::terminate::Terminate; @@ -30,7 +33,7 @@ use crate::proxy::EncryptionService; use crate::{EqlOutput, EqlQueryPayload}; use bytes::BytesMut; use cipherstash_client::encryption::Plaintext; -use eql_mapper::{self, EqlMapperError, EqlTerm, TypeCheckedStatement}; +use eql_mapper::{self, EqlMapperError, EqlTermVariant, JsonSelectorSource, TypeCheckedStatement}; use metrics::{counter, histogram}; use pg_escape::quote_literal; use serde::Serialize; @@ -598,7 +601,7 @@ where return Ok(vec![]); } - let plaintexts = literals_to_plaintext(literal_values, literal_columns)?; + let plaintexts = literals_to_plaintext(typed_statement, literal_columns)?; let start = Instant::now(); @@ -816,6 +819,7 @@ where counter!(STATEMENTS_ENCRYPTED_TOTAL).increment(1); message.rewrite_param_types(&statement.param_columns); + message.declare_unreferenced_param_types(&statement.unreferenced_param_indexes()); self.context .add_statement(message.name.to_owned(), statement); } @@ -931,6 +935,7 @@ where projection_columns.to_owned(), literal_columns.to_owned(), param_types, + json_value_selector_params(typed_statement), ); Ok(Some(statement)) @@ -1051,8 +1056,11 @@ where bind: &Bind, statement: &Statement, ) -> Result>, Error> { - let plaintexts = - bind.to_plaintext(&statement.param_columns, &statement.postgres_param_types)?; + let plaintexts = bind.to_plaintext( + &statement.param_columns, + &statement.postgres_param_types, + &statement.json_value_selectors, + )?; debug!(target: MAPPER, client_id = self.context.client_id, plaintexts = ?plaintexts); @@ -1166,30 +1174,94 @@ where } } +/// The fused JSON value-selector params of a statement, keyed by 0-based bind +/// index. [`eql_mapper::Param`] is 1-based, so every index shifts by one here. +/// +/// A value-selector param whose path is *also* a param records that param's +/// index; a literal path is carried inline. +fn json_value_selector_params( + typed_statement: &TypeCheckedStatement<'_>, +) -> HashMap { + typed_statement + .params + .iter() + .enumerate() + .filter_map(|(idx, (param, _))| { + let path = match typed_statement.json_value_selectors.for_param(*param)? { + JsonSelectorSource::Literal(path) => JsonSelectorPath::Literal(path.to_owned()), + JsonSelectorSource::Param(path_param) => { + JsonSelectorPath::Param(path_param.0.saturating_sub(1) as usize) + } + }; + Some((idx, path)) + }) + .collect() +} + fn literals_to_plaintext( - literals: &Vec<(EqlTerm, &ast::Value)>, + typed_statement: &TypeCheckedStatement<'_>, literal_columns: &Vec>, ) -> Result>, Error> { + let literals = typed_statement.literal_values(); + let plaintexts = literals .iter() .zip(literal_columns) - .map(|((_, val), col)| match col { - Some(col) => literal_from_sql(val, col.eql_term(), col.cast_type()).map_err(|err| { - debug!( - target: MAPPER, - msg = "Could not convert literal value", - value = ?val, - cast_type = ?col.cast_type(), - error = err.to_string() - ); - MappingError::InvalidParameter(Box::new(col.to_owned())) - }), + .map(|((eql_term, val), col)| match col { + Some(col) => { + let plaintext = if eql_term.variant() == EqlTermVariant::JsonValueSelector { + json_value_selector_literal_plaintext(typed_statement, val) + } else { + literal_from_sql(val, col.eql_term(), col.cast_type()) + }; + + plaintext.map_err(|err| { + debug!( + target: MAPPER, + msg = "Could not convert literal value", + value = ?val, + cast_type = ?col.cast_type(), + error = err.to_string() + ); + MappingError::InvalidParameter(Box::new(col.to_owned())).into() + }) + } None => Ok(None), }) - .collect::, _>>()?; + .collect::, Error>>()?; Ok(plaintexts) } +/// Composes the needle for a JSON field equality whose value is a literal: +/// `{"path": , "value": }`. +/// +/// Only a literal path can be resolved here — the whole statement is encrypted +/// at Parse time, before any param is bound. `col -> $1 = 'value'` (param path, +/// literal value) is therefore not supported; it is also not a shape any client +/// produces, since a client that parameterises the path parameterises the value +/// too. +fn json_value_selector_literal_plaintext( + typed_statement: &TypeCheckedStatement<'_>, + literal: &ast::Value, +) -> Result, MappingError> { + let Some(JsonSelectorSource::Literal(path)) = + typed_statement.json_value_selectors.for_literal(literal) + else { + debug!( + target: MAPPER, + msg = "Encrypted JSON equality needs a literal selector when the value is a literal", + value = ?literal, + ); + return Err(MappingError::CouldNotParseParameter); + }; + + let Some(value) = literal_json_value(literal)? else { + return Ok(None); + }; + + json_value_selector_plaintext(path, value).map(Some) +} + fn to_json_literal_value(literal: &T) -> Result where T: ?Sized + Serialize, diff --git a/packages/cipherstash-proxy/src/postgresql/messages/bind.rs b/packages/cipherstash-proxy/src/postgresql/messages/bind.rs index 54b538214..117d94834 100644 --- a/packages/cipherstash-proxy/src/postgresql/messages/bind.rs +++ b/packages/cipherstash-proxy/src/postgresql/messages/bind.rs @@ -2,7 +2,10 @@ use super::{maybe_json, maybe_jsonb, Name, NULL}; use crate::error::{Error, MappingError, ProtocolError}; use crate::log::MAPPER; use crate::postgresql::context::column::Column; -use crate::postgresql::data::bind_param_from_sql; +use crate::postgresql::context::statement::JsonSelectorPath; +use crate::postgresql::data::{ + bind_param_from_sql, bind_param_json_value, json_value_selector_plaintext, +}; use crate::postgresql::format_code::FormatCode; use crate::postgresql::protocol::BytesMutReadString; use crate::{EqlOutput, EqlQueryPayload}; @@ -10,6 +13,7 @@ use crate::{SIZE_I16, SIZE_I32}; use bytes::{Buf, BufMut, BytesMut}; use cipherstash_client::encryption::Plaintext; use postgres_types::Type; +use std::collections::HashMap; use std::fmt::{self, Display, Formatter}; use std::io::Cursor; use std::{convert::TryFrom, ffi::CString}; @@ -44,10 +48,89 @@ impl Bind { .any(|param| param.requires_rewrite()) } + /// Converts the bound params to plaintexts, one per param. + /// + /// `json_value_selectors` names the params that are *fused*: a JSON field + /// equality (`col -> $1 = $2`) has no plaintext of its own for `$2` — its + /// needle is composed from `$1` and `$2` together. Those params are resolved + /// in a second pass, once every param's own value is available. pub fn to_plaintext( &self, param_columns: &[Option], param_types: &[i32], + json_value_selectors: &HashMap, + ) -> Result>, Error> { + let mut plaintexts = + self.to_plaintext_positional(param_columns, param_types, json_value_selectors)?; + + for (idx, path) in json_value_selectors.iter() { + if *idx >= plaintexts.len() { + continue; + } + + let postgres_type = param_columns + .get(*idx) + .and_then(|col| col.as_ref()) + .map(|col| get_param_type(*idx, param_types, col)) + .unwrap_or(Type::JSONB); + + plaintexts[*idx] = self.json_value_selector_plaintext(*idx, path, &postgres_type)?; + } + + Ok(plaintexts) + } + + /// Composes `{"path", "value"}` for the fused value-selector param at `idx`. + /// + /// The path comes either from the SQL (a literal selector) or from another + /// bind param, which is read from the wire directly rather than from its + /// plaintext — the plaintext pass has already been run for that param, but + /// its value is the selector *text*, and reading the raw bytes keeps the two + /// halves decoded the same way regardless of which pass ran first. + fn json_value_selector_plaintext( + &self, + idx: usize, + path: &JsonSelectorPath, + postgres_type: &Type, + ) -> Result, Error> { + let path = match path { + JsonSelectorPath::Literal(path) => path.to_owned(), + JsonSelectorPath::Param(path_idx) => match self.param_values.get(*path_idx) { + Some(param) if !param.is_null() => param.to_string(), + _ => return Ok(None), + }, + }; + + let Some(param) = self.param_values.get(idx) else { + return Ok(None); + }; + + let Some(value) = bind_param_json_value(param, postgres_type)? else { + return Ok(None); + }; + + debug!( + target: MAPPER, + msg = "Fused JSON value selector", + ?idx, + ?path, + ?value + ); + + Ok(Some(json_value_selector_plaintext(&path, value)?)) + } + + /// The one-plaintext-per-param pass. + /// + /// Params named in `json_value_selectors` are skipped: a fused value + /// selector has no plaintext of its own — its own bytes are only half the + /// needle — and decoding it here would fail, since the scalar it carries is + /// not a valid standalone operand for the column. + fn to_plaintext_positional( + &self, + param_columns: &[Option], + param_types: &[i32], + json_value_selectors: &HashMap, ) -> Result>, Error> { let plaintexts = self .param_values @@ -55,6 +138,7 @@ impl Bind { .zip(param_columns.iter()) .enumerate() .map(|(idx, (param, col))| match col { + _ if json_value_selectors.contains_key(&idx) => Ok(None), Some(col) => { let bound_param_type = get_param_type(idx, param_types, col); diff --git a/packages/cipherstash-proxy/src/postgresql/messages/parse.rs b/packages/cipherstash-proxy/src/postgresql/messages/parse.rs index 10a2038aa..f89073ed4 100644 --- a/packages/cipherstash-proxy/src/postgresql/messages/parse.rs +++ b/packages/cipherstash-proxy/src/postgresql/messages/parse.rs @@ -8,6 +8,9 @@ use bytes::{Buf, BufMut, BytesMut}; use postgres_types::Type; use std::{ffi::CString, io::Cursor}; +/// PostgreSQL's "unspecified type, infer it" param OID in a Parse message. +const UNSPECIFIED_TYPE_OID: i32 = 0; + #[derive(Debug, Clone)] pub struct Parse { pub code: char, @@ -40,6 +43,39 @@ impl Parse { } } + /// + /// Declares a param type for every placeholder the rewrite left + /// unreferenced. + /// + /// A JSON field equality (`col -> $1 = $2`) fuses both operands into one + /// needle, so `$1` disappears from the rewritten SQL while staying in the + /// client's Bind. PostgreSQL allows an unreferenced param, but only if its + /// type is declared — otherwise it fails with "could not determine data type + /// of parameter $1". Clients that let the server infer types (the common + /// case) send no types at all, so the array is grown here. + /// + /// Slots that are still referenced are left as OID 0, which PostgreSQL reads + /// as "unspecified, infer it" — so this never overrides a type the query + /// itself determines. The dropped selectors are bound as the encrypted + /// selector *text*, hence TEXT. + /// + pub fn declare_unreferenced_param_types(&mut self, indexes: &[usize]) { + let Some(required_len) = indexes.iter().max().map(|idx| idx + 1) else { + return; + }; + + if self.param_types.len() < required_len { + self.param_types.resize(required_len, UNSPECIFIED_TYPE_OID); + self.num_params = self.param_types.len() as i16; + self.dirty = true; + } + + for idx in indexes { + self.param_types[*idx] = Type::TEXT.oid() as i32; + self.dirty = true; + } + } + pub fn rewrite_statement(&mut self, statement: String) { self.statement = statement; self.dirty = true; diff --git a/packages/cipherstash-proxy/src/proxy/zerokms/zerokms.rs b/packages/cipherstash-proxy/src/proxy/zerokms/zerokms.rs index 61f837663..d2cf1d8a6 100644 --- a/packages/cipherstash-proxy/src/proxy/zerokms/zerokms.rs +++ b/packages/cipherstash-proxy/src/proxy/zerokms/zerokms.rs @@ -283,6 +283,23 @@ impl EncryptionService for ZeroKms { .find(|i| matches!(i.index_type, IndexType::SteVec { .. })) .map(|index| EqlOperation::Query(&index.index_type, QueryOp::SteVecTerm)) .unwrap_or(EqlOperation::Store), + + // JsonValueSelector is the fused value operand of a JSON + // field equality (`col -> sel = value`). Its plaintext is the + // composition input `{"path", "value"}` (built by the + // frontend from BOTH SQL operands); the client MACs them + // together into one selector, applying the column's term + // filters to the value. The result is a one-entry containment + // needle matched by `eql_v3.jsonb_contains`. + EqlTermVariant::JsonValueSelector => col + .config + .indexes + .iter() + .find(|i| matches!(i.index_type, IndexType::SteVec { .. })) + .map(|index| { + EqlOperation::Query(&index.index_type, QueryOp::SteVecValueSelector) + }) + .unwrap_or(EqlOperation::Store), }; let prepared = PreparedPlaintext::new( From 15066cca6ef2d79d2c261eda1c45623a748a7c89 Mon Sep 17 00:00:00 2001 From: James Sadler Date: Sun, 26 Jul 2026 01:10:37 +1000 Subject: [PATCH 03/10] docs: changelog entry for encrypted JSON field equality Stable-Commit-Id: q-6658b8n0t0qbfkcksk6720x0b5 --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 134101a05..1d1f428a6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). - **Encrypted full-text match with `@@`**: The `@@` operator is now supported on encrypted text columns whose domain carries a match (bloom-filter) term, rewritten to the EQL v3 `eql_v3.match_term` form. +- **Equality on encrypted JSON fields**: `WHERE col -> 'field' = 'value'` now works on encrypted JSON columns, in both the simple and extended query protocols, and in the `->>` and `jsonb_path_query_first(col, path) = value` spellings. `<>` is supported as the negation. The field and the value are combined into a single encrypted value-selector needle and matched by containment, so a query never reveals the field and value separately. Matching is exact and case-sensitive; the value must be a JSON scalar (comparing a whole object or array to a field is rejected — use containment with `@>` instead). + ### Fixed - **`LIKE`/`ILIKE` capability checking**: `LIKE` and `ILIKE` on an encrypted column are now gated by the column's token-match capability. Previously these predicates bypassed capability checking and were silently accepted on columns that do not support fuzzy match; they are now rejected with a capability error. From b7c7a7762e29db314d130779ab4042e4e0ad3bba Mon Sep 17 00:00:00 2001 From: James Sadler Date: Sun, 26 Jul 2026 10:40:58 +1000 Subject: [PATCH 04/10] refactor: model output params as derived from N input params MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the dangling-param trick with the real shape of the problem: an output param may be derived from more than one input param, so there is no 1:1 correspondence to rely on. Previously the JSON-equality rewrite left the fused path operand in the SQL as an unreferenced `$n`, purely so that param numbering stayed positional. That kept the plumbing simple at the cost of lying about the statement: a param that means nothing, and a neighbouring param whose meaning has quietly become "composite". Both are invisible at the call sites that bind them. The correspondence is now explicit: - `TypeCheckedStatement::transform` returns a `TransformedStatement` — the rewritten SQL plus a `ParamPlan`. Each `OutputParam` names the input param(s) its value is built from (`Input`, or `JsonValueSelector` carrying both operands). - A `RenumberParams` pass runs after the rewrite rules and assigns `$1..$m` in SQL order, so rules may freely drop or duplicate placeholders without maintaining numbering themselves. Running it as a separate pass keeps `FailOnPlaceholderChange` governing the rules. - `ParamPlan::check_covers` enforces the invariant that replaces 1:1 — **coverage**: every input param must be consumed by some output. An input that feeds nothing has been silently dropped, which is a bug in a rewrite rule. The proxy binds against the plan rather than by position. `Statement` carries both shapes: `param_columns` (what the client binds, used to decode and to answer Describe) and `output_params` (what PostgreSQL receives). Bind builds each output from the inputs its source names; ParameterDescription is rebuilt from the input params, since the server describes the rewritten statement and would otherwise tell the client too few params to bind. Parse carries each client-declared type across to the output param that consumes it. Every output param is now referenced by the rewritten SQL, so `declare_unreferenced_param_types` and its unreferenced-placeholder workaround are gone. When the plan is positional — every statement that is not a JSON equality — Bind still patches values in place, leaving the client's wire framing byte-for-byte as sent. Mapper units 102 -> 104 (renumbering across a fusion, and identity plans); proxy units 116 -> 117. Integration unchanged: the 2 JSON equality tests pass, and the four ORE-ordering tests that differ from the previous run fail identically on the base binary when run in isolation. Stable-Commit-Id: q-67avvp278knhdy7xe192mq667s --- packages/cipherstash-proxy/src/error.rs | 3 + .../src/postgresql/backend.rs | 26 +- .../src/postgresql/column_mapper.rs | 36 ++- .../src/postgresql/context/mod.rs | 15 +- .../src/postgresql/context/statement.rs | 132 ++++++++-- .../src/postgresql/frontend.rs | 88 +++---- .../src/postgresql/messages/bind.rs | 230 ++++++++++-------- .../src/postgresql/messages/mod.rs | 4 + .../postgresql/messages/param_description.rs | 15 ++ .../src/postgresql/messages/parse.rs | 129 ++++++---- packages/eql-mapper/src/lib.rs | 81 +++++- packages/eql-mapper/src/param_plan.rs | 120 +++++++++ packages/eql-mapper/src/renumber_params.rs | 60 +++++ .../eql-mapper/src/type_checked_statement.rs | 105 +++++++- 14 files changed, 792 insertions(+), 252 deletions(-) create mode 100644 packages/eql-mapper/src/param_plan.rs create mode 100644 packages/eql-mapper/src/renumber_params.rs diff --git a/packages/cipherstash-proxy/src/error.rs b/packages/cipherstash-proxy/src/error.rs index 1068e66cf..fb3982b4f 100644 --- a/packages/cipherstash-proxy/src/error.rs +++ b/packages/cipherstash-proxy/src/error.rs @@ -440,6 +440,9 @@ pub enum ProtocolError { #[error("Expected {expected} parameter format codes, received {received}")] ParameterResultFormatCodesMismatch { expected: usize, received: usize }, + #[error("Rewritten statement binds parameter {param}, but only {received} were provided")] + MissingBoundParameter { param: usize, received: usize }, + #[error("Expected a {expected} message, received message code {received}")] UnexpectedAuthenticationResponse { expected: String, received: i32 }, diff --git a/packages/cipherstash-proxy/src/postgresql/backend.rs b/packages/cipherstash-proxy/src/postgresql/backend.rs index 616f456cb..d22730092 100644 --- a/packages/cipherstash-proxy/src/postgresql/backend.rs +++ b/packages/cipherstash-proxy/src/postgresql/backend.rs @@ -4,7 +4,7 @@ use super::error_handler::PostgreSqlErrorHandler; use super::message_buffer::MessageBuffer; use super::messages::error_response::ErrorResponse; use super::messages::row_description::RowDescription; -use super::messages::BackendCode; +use super::messages::{BackendCode, UNSPECIFIED_TYPE_OID}; use super::Column; use crate::connect::Sender; use crate::error::{EncryptError, Error}; @@ -572,20 +572,34 @@ where debug!(target: PROTOCOL, client_id = self.context.client_id, ParamDescription = ?description); if let Some(statement) = self.context.get_statement_from_describe() { + // Describe the params the CLIENT wrote, not the ones PostgreSQL was + // sent. A rewrite may have fused or dropped params, in which case + // the server's description is both shorter than and shifted from + // what the client needs in order to bind. let param_types = statement .param_columns .iter() - .map(|col| { - col.as_ref().map(|col| { + .enumerate() + .map(|(idx, col)| match col { + Some(col) => { debug!(target: MAPPER, client_id = self.context.client_id, ColumnConfig = ?col); - col.postgres_type.clone() - }) + col.postgres_type.oid() as i32 + } + // A native param is never fused, so it reaches PostgreSQL + // as some output param; take the type the server inferred + // for it. + None => statement + .output_params + .iter() + .position(|output| output.source.primary_input() == idx) + .and_then(|output_idx| description.types.get(output_idx).copied()) + .unwrap_or(UNSPECIFIED_TYPE_OID), }) .collect::>(); debug!(target: MAPPER, client_id = self.context.client_id, param_types = ?param_types); - description.map_types(¶m_types); + description.set_types(param_types); } if description.requires_rewrite() { diff --git a/packages/cipherstash-proxy/src/postgresql/column_mapper.rs b/packages/cipherstash-proxy/src/postgresql/column_mapper.rs index b8600e318..2e1f93f1c 100644 --- a/packages/cipherstash-proxy/src/postgresql/column_mapper.rs +++ b/packages/cipherstash-proxy/src/postgresql/column_mapper.rs @@ -5,7 +5,7 @@ use crate::{ proxy::EncryptConfig, }; use cipherstash_client::eql::Identifier; -use eql_mapper::{EqlTerm, TableColumn, TypeCheckedStatement}; +use eql_mapper::{EqlTerm, ParamPlan, TableColumn, TypeCheckedStatement}; use postgres_types::Type; use std::sync::Arc; use tracing::{debug, warn}; @@ -95,6 +95,40 @@ impl ColumnMapper { Ok(param_columns) } + /// Maps the params of the *rewritten* statement to an Encrypt column + /// configuration, positionally over [`ParamPlan::outputs`]. + /// + /// These are the values actually sent to PostgreSQL, which after a fusion + /// are not the values the client bound — hence a separate mapping from + /// [`Self::get_param_columns`]. + pub fn get_output_param_columns(&self, plan: &ParamPlan) -> Result>, Error> { + let mut output_columns = vec![]; + + for output in plan.outputs() { + let configured_column = match &output.value { + eql_mapper::Value::Eql(eql_term) => { + let TableColumn { table, column } = eql_term.table_column(); + let identifier = + Identifier::new(table.value.to_string(), column.value.to_string()); + + debug!( + target: MAPPER, + msg = "Encrypted output parameter", + param = %output.param, + column = ?identifier, + ?eql_term, + ); + + self.get_column(identifier, eql_term)? + } + _ => None, + }; + output_columns.push(configured_column); + } + + Ok(output_columns) + } + /// Maps typed statement literal columns to an Encrypt column configuration pub fn get_literal_columns( &self, diff --git a/packages/cipherstash-proxy/src/postgresql/context/mod.rs b/packages/cipherstash-proxy/src/postgresql/context/mod.rs index 50411dc47..d42e015cf 100644 --- a/packages/cipherstash-proxy/src/postgresql/context/mod.rs +++ b/packages/cipherstash-proxy/src/postgresql/context/mod.rs @@ -3,11 +3,7 @@ pub mod phase_timing; pub mod portal; pub mod statement; pub mod statement_metadata; -pub use self::{ - phase_timing::PhaseTiming, - portal::Portal, - statement::{JsonSelectorPath, Statement}, -}; +pub use self::{phase_timing::PhaseTiming, portal::Portal, statement::Statement}; use super::{ column_mapper::ColumnMapper, messages::{describe::Describe, Name, Target}, @@ -818,6 +814,13 @@ where self.column_mapper.get_param_columns(typed_statement) } + pub fn get_output_param_columns( + &self, + plan: &eql_mapper::ParamPlan, + ) -> Result>, Error> { + self.column_mapper.get_output_param_columns(plan) + } + pub fn get_literal_columns( &self, typed_statement: &eql_mapper::TypeCheckedStatement<'_>, @@ -1120,7 +1123,7 @@ mod tests { projection_columns: vec![], literal_columns: vec![], postgres_param_types: vec![], - json_value_selectors: std::collections::HashMap::new(), + output_params: vec![], } } diff --git a/packages/cipherstash-proxy/src/postgresql/context/statement.rs b/packages/cipherstash-proxy/src/postgresql/context/statement.rs index 6bd90d6df..920897361 100644 --- a/packages/cipherstash-proxy/src/postgresql/context/statement.rs +++ b/packages/cipherstash-proxy/src/postgresql/context/statement.rs @@ -1,8 +1,7 @@ use super::Column; -use std::collections::HashMap; +use eql_mapper::{JsonSelectorSource, ParamPlan}; -/// Where the path half of a fused JSON value selector comes from, resolved to -/// this statement's bind params. +/// Where the path half of a fused JSON value selector comes from. /// /// The proxy's copy of [`eql_mapper::JsonSelectorSource`], with param numbers /// converted to 0-based bind indexes. @@ -11,59 +10,86 @@ pub enum JsonSelectorPath { /// A literal path in the SQL, known at Parse time. Literal(String), - /// A placeholder path, arriving in this (0-based) bind param. + /// A placeholder path, arriving in this (0-based) input param. Param(usize), } +/// How the value bound to one output param is built from the input params. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum OutputParamSource { + /// Taken from this (0-based) input param. + Input(usize), + + /// Fused from a JSON path and a value into one value-selector needle. + JsonValueSelector { + path: JsonSelectorPath, + value: usize, + }, +} + +impl OutputParamSource { + /// The input param this output is built *around* — the one whose wire + /// format and (for a passthrough) whose bytes it inherits. For a fusion + /// that is the value operand; the path only contributes to the needle. + pub fn primary_input(&self) -> usize { + match self { + OutputParamSource::Input(idx) => *idx, + OutputParamSource::JsonValueSelector { value, .. } => *value, + } + } +} + +/// One param of the *rewritten* statement — what PostgreSQL will be sent. +#[derive(Debug, Clone, PartialEq)] +pub struct OutputParam { + /// The column configuration, when this param must be encrypted. `None` for + /// a native param, which is forwarded byte-for-byte. + pub column: Option, + + /// The input param(s) its value is built from. + pub source: OutputParamSource, +} + /// /// Type Analysed parameters and projection /// +/// Params have **two** shapes, and they are not guaranteed to correspond: +/// `param_columns` describes what the *client* binds, `output_params` describes +/// what PostgreSQL receives. Encrypted JSON equality fuses two input params into +/// one output param, so any code that assumes `$n` in equals `$n` out is wrong. +/// #[derive(Debug, Clone, PartialEq)] pub struct Statement { + /// The params as the client sees them — used to decode bound values and to + /// answer `Describe`. pub param_columns: Vec>, + + /// The params of the rewritten statement, in the order PostgreSQL sees + /// them, each naming the input it is built from. + pub output_params: Vec, + pub projection_columns: Vec>, pub literal_columns: Vec>, pub postgres_param_types: Vec, - - /// Params that carry a fused JSON value-selector needle (`col -> sel = - /// $n`), keyed by 0-based bind index, mapped to where their path comes from. - /// - /// Bind consults this to compose `{"path", "value"}` from two params before - /// encrypting. Empty for every statement without encrypted JSON equality. - pub json_value_selectors: HashMap, } impl Statement { pub fn new( param_columns: Vec>, + output_params: Vec, projection_columns: Vec>, literal_columns: Vec>, postgres_param_types: Vec, - json_value_selectors: HashMap, ) -> Statement { Statement { param_columns, + output_params, projection_columns, literal_columns, postgres_param_types, - json_value_selectors, } } - /// The 0-based bind indexes of params that supply only a value-selector - /// path. The rewrite folds these into the needle and drops them from the - /// SQL, so PostgreSQL never references them — but it still needs their type - /// declared in Parse to prepare the statement. - pub fn unreferenced_param_indexes(&self) -> Vec { - self.json_value_selectors - .values() - .filter_map(|path| match path { - JsonSelectorPath::Param(idx) => Some(*idx), - JsonSelectorPath::Literal(_) => None, - }) - .collect() - } - pub fn has_literals(&self) -> bool { !self.literal_columns.is_empty() } @@ -76,3 +102,55 @@ impl Statement { !self.projection_columns.is_empty() } } + +/// `true` when `output_params` are the first `output_params.len()` input params +/// in order, unchanged — the ordinary case, where the rewrite reshaped nothing. +/// +/// Callers that hold the bound values must also check the count matches: a +/// prefix match alone would silently drop trailing params. +pub fn params_are_positional(output_params: &[OutputParam]) -> bool { + output_params + .iter() + .enumerate() + .all(|(idx, output)| output.source == OutputParamSource::Input(idx)) +} + +/// Converts a mapper [`ParamPlan`] to the proxy's 0-based form, pairing each +/// output param with the column configuration that says how to encrypt it. +/// +/// `output_columns` is positional over the plan's outputs. +pub fn output_params_from_plan( + plan: &ParamPlan, + output_columns: Vec>, +) -> Vec { + plan.outputs() + .iter() + .zip(output_columns) + .map(|(output, column)| OutputParam { + column, + source: match &output.source { + eql_mapper::OutputParamSource::Input(param) => { + OutputParamSource::Input(to_index(param.0)) + } + eql_mapper::OutputParamSource::JsonValueSelector { path, value } => { + OutputParamSource::JsonValueSelector { + path: match path { + JsonSelectorSource::Literal(path) => { + JsonSelectorPath::Literal(path.to_owned()) + } + JsonSelectorSource::Param(param) => { + JsonSelectorPath::Param(to_index(param.0)) + } + }, + value: to_index(value.0), + } + } + }, + }) + .collect() +} + +/// Mapper params are 1-based (`$1`), bind params are 0-based. +fn to_index(param: u16) -> usize { + param.saturating_sub(1) as usize +} diff --git a/packages/cipherstash-proxy/src/postgresql/frontend.rs b/packages/cipherstash-proxy/src/postgresql/frontend.rs index e6d328516..3de073628 100644 --- a/packages/cipherstash-proxy/src/postgresql/frontend.rs +++ b/packages/cipherstash-proxy/src/postgresql/frontend.rs @@ -13,8 +13,10 @@ use crate::connect::Sender; use crate::error::{EncryptError, Error, MappingError}; use crate::log::{MAPPER, PROTOCOL}; use crate::postgresql::context::column::Column; +use crate::postgresql::context::statement::{ + output_params_from_plan, OutputParam, OutputParamSource, +}; use crate::postgresql::context::statement_metadata::{ProtocolType, StatementType}; -use crate::postgresql::context::JsonSelectorPath; use crate::postgresql::context::Portal; use crate::postgresql::data::{ json_value_selector_plaintext, literal_from_sql, literal_json_value, @@ -469,10 +471,12 @@ where { debug!(target: MAPPER, client_id = self.context.client_id, - transformed_statement = ?transformed_statement, + transformed_statement = ?transformed_statement.statement, ); - transformed_statements.push(transformed_statement); + // The simple protocol has no params, so the plan is + // always empty here — only the SQL is needed. + transformed_statements.push(transformed_statement.statement); encrypted = true; } } @@ -647,7 +651,7 @@ where &mut self, typed_statement: &TypeCheckedStatement<'_>, encrypted_literals: &Vec>, - ) -> Result, Error> { + ) -> Result, Error> { // Convert literals to ast Expr let mut encrypted_expressions = vec![]; for encrypted in encrypted_literals { @@ -792,7 +796,7 @@ where let mut parse_duration_recorded = false; match self.to_encryptable_statement(&typed_statement, param_types)? { - Some(statement) => { + Some(mut statement) => { if typed_statement.requires_transform() { // Record parse duration before encryption work starts self.context @@ -809,17 +813,26 @@ where { debug!(target: MAPPER, client_id = self.context.client_id, - transformed_statement = ?transformed_statement, + transformed_statement = ?transformed_statement.statement, + param_plan = ?transformed_statement.params, ); - message.rewrite_statement(transformed_statement.to_string()); + // The rewrite may have reshaped the params, so the + // statement's output params come from the plan rather + // than from its own input params. + let output_columns = self + .context + .get_output_param_columns(&transformed_statement.params)?; + statement.output_params = + output_params_from_plan(&transformed_statement.params, output_columns); + + message.rewrite_statement(transformed_statement.statement.to_string()); } } counter!(STATEMENTS_ENCRYPTED_TOTAL).increment(1); - message.rewrite_param_types(&statement.param_columns); - message.declare_unreferenced_param_types(&statement.unreferenced_param_indexes()); + message.rewrite_param_types(&statement.output_params); self.context .add_statement(message.name.to_owned(), statement); } @@ -930,12 +943,24 @@ where literal_columns = ?literal_columns, ); + // Until the statement is rewritten its output params are its input + // params — a statement that needs no transform never reshapes them, and + // one that does overwrites this from the rewrite's `ParamPlan`. + let output_params = param_columns + .iter() + .enumerate() + .map(|(idx, column)| OutputParam { + column: column.to_owned(), + source: OutputParamSource::Input(idx), + }) + .collect(); + let statement = Statement::new( param_columns.to_owned(), + output_params, projection_columns.to_owned(), literal_columns.to_owned(), param_types, - json_value_selector_params(typed_statement), ); Ok(Some(statement)) @@ -1012,7 +1037,7 @@ where if statement.has_params() { let encrypted = self.encrypt_params(session_id, &bind, &statement).await?; - bind.rewrite(encrypted)?; + bind.rewrite(&statement.output_params, encrypted)?; } if statement.has_projection() { portal = Portal::encrypted_with_format_codes( @@ -1056,11 +1081,16 @@ where bind: &Bind, statement: &Statement, ) -> Result>, Error> { - let plaintexts = bind.to_plaintext( - &statement.param_columns, - &statement.postgres_param_types, - &statement.json_value_selectors, - )?; + let plaintexts = + bind.to_plaintext(&statement.output_params, &statement.postgres_param_types)?; + + // Encryption is positional over the OUTPUT params — the values actually + // sent — not over what the client bound. + let output_param_columns = statement + .output_params + .iter() + .map(|output| output.column.to_owned()) + .collect::>(); debug!(target: MAPPER, client_id = self.context.client_id, plaintexts = ?plaintexts); @@ -1068,7 +1098,7 @@ where let encrypted = self .context - .encrypt(plaintexts, &statement.param_columns) + .encrypt(plaintexts, &output_param_columns) .await .inspect_err(|_| { counter!(ENCRYPTION_ERROR_TOTAL).increment(1); @@ -1174,30 +1204,6 @@ where } } -/// The fused JSON value-selector params of a statement, keyed by 0-based bind -/// index. [`eql_mapper::Param`] is 1-based, so every index shifts by one here. -/// -/// A value-selector param whose path is *also* a param records that param's -/// index; a literal path is carried inline. -fn json_value_selector_params( - typed_statement: &TypeCheckedStatement<'_>, -) -> HashMap { - typed_statement - .params - .iter() - .enumerate() - .filter_map(|(idx, (param, _))| { - let path = match typed_statement.json_value_selectors.for_param(*param)? { - JsonSelectorSource::Literal(path) => JsonSelectorPath::Literal(path.to_owned()), - JsonSelectorSource::Param(path_param) => { - JsonSelectorPath::Param(path_param.0.saturating_sub(1) as usize) - } - }; - Some((idx, path)) - }) - .collect() -} - fn literals_to_plaintext( typed_statement: &TypeCheckedStatement<'_>, literal_columns: &Vec>, diff --git a/packages/cipherstash-proxy/src/postgresql/messages/bind.rs b/packages/cipherstash-proxy/src/postgresql/messages/bind.rs index 117d94834..410436953 100644 --- a/packages/cipherstash-proxy/src/postgresql/messages/bind.rs +++ b/packages/cipherstash-proxy/src/postgresql/messages/bind.rs @@ -2,7 +2,9 @@ use super::{maybe_json, maybe_jsonb, Name, NULL}; use crate::error::{Error, MappingError, ProtocolError}; use crate::log::MAPPER; use crate::postgresql::context::column::Column; -use crate::postgresql::context::statement::JsonSelectorPath; +use crate::postgresql::context::statement::{ + params_are_positional, JsonSelectorPath, OutputParam, OutputParamSource, +}; use crate::postgresql::data::{ bind_param_from_sql, bind_param_json_value, json_value_selector_plaintext, }; @@ -13,7 +15,6 @@ use crate::{SIZE_I16, SIZE_I32}; use bytes::{Buf, BufMut, BytesMut}; use cipherstash_client::encryption::Plaintext; use postgres_types::Type; -use std::collections::HashMap; use std::fmt::{self, Display, Formatter}; use std::io::Cursor; use std::{convert::TryFrom, ffi::CString}; @@ -32,6 +33,10 @@ pub struct Bind { pub param_values: Vec, pub num_result_column_format_codes: i16, pub result_columns_format_codes: Vec, + /// Set when the param list was rebuilt because the rewrite reshaped the + /// params. The message must then be re-sent even if no individual param was + /// itself edited, because the count and framing changed. + reshaped: bool, } #[derive(Clone, Debug, PartialEq)] @@ -43,54 +48,78 @@ pub struct BindParam { impl Bind { pub fn requires_rewrite(&self) -> bool { - self.param_values - .iter() - .any(|param| param.requires_rewrite()) + self.reshaped + || self + .param_values + .iter() + .any(|param| param.requires_rewrite()) } - /// Converts the bound params to plaintexts, one per param. + /// Converts the bound params to the plaintexts of the **output** params — + /// the values PostgreSQL will receive, which are not necessarily the values + /// the client bound. /// - /// `json_value_selectors` names the params that are *fused*: a JSON field - /// equality (`col -> $1 = $2`) has no plaintext of its own for `$2` — its - /// needle is composed from `$1` and `$2` together. Those params are resolved - /// in a second pass, once every param's own value is available. + /// Each output param pulls from the input param(s) its `source` names, so a + /// fused JSON value selector reads both halves here and a dropped path + /// operand is never decoded on its own (its bytes are only half a needle and + /// would not decode as a standalone operand for the column). pub fn to_plaintext( &self, - param_columns: &[Option], + output_params: &[OutputParam], param_types: &[i32], - json_value_selectors: &HashMap, ) -> Result>, Error> { - let mut plaintexts = - self.to_plaintext_positional(param_columns, param_types, json_value_selectors)?; - - for (idx, path) in json_value_selectors.iter() { - if *idx >= plaintexts.len() { - continue; - } - - let postgres_type = param_columns - .get(*idx) - .and_then(|col| col.as_ref()) - .map(|col| get_param_type(*idx, param_types, col)) - .unwrap_or(Type::JSONB); - - plaintexts[*idx] = self.json_value_selector_plaintext(*idx, path, &postgres_type)?; - } - - Ok(plaintexts) + output_params + .iter() + .map(|output| { + let Some(col) = &output.column else { + // Native param: forwarded verbatim, nothing to encrypt. + return Ok(None); + }; + + let input = output.source.primary_input(); + let bound_param_type = get_param_type(input, param_types, col); + + debug!( + target: MAPPER, + col = ?col, bound_param_type = ?bound_param_type, ?input + ); + + match &output.source { + OutputParamSource::Input(idx) => { + let Some(param) = self.param_values.get(*idx) else { + return Ok(None); + }; + + // Convert param bytes into a Plaintext wrapping a Value + // If the param type is different, will convert the bound type to the correct Plaintext variant identified by the cast_type + bind_param_from_sql( + param, + &bound_param_type, + col.eql_term(), + col.cast_type(), + ) + .map_err(|_| { + MappingError::InvalidParameter(Box::new(col.to_owned())).into() + }) + } + OutputParamSource::JsonValueSelector { path, value } => { + self.json_value_selector_plaintext(path, *value, &bound_param_type) + } + } + }) + .collect() } - /// Composes `{"path", "value"}` for the fused value-selector param at `idx`. + /// Composes `{"path", "value"}` — the input to `SteVecValueSelector` — from + /// the two operands of a JSON field equality. /// - /// The path comes either from the SQL (a literal selector) or from another - /// bind param, which is read from the wire directly rather than from its - /// plaintext — the plaintext pass has already been run for that param, but - /// its value is the selector *text*, and reading the raw bytes keeps the two - /// halves decoded the same way regardless of which pass ran first. + /// The path is either a literal from the SQL or another bind param, which is + /// read straight off the wire: it is the selector *text*, so it needs none + /// of the per-column decoding the value half goes through. fn json_value_selector_plaintext( &self, - idx: usize, path: &JsonSelectorPath, + value: usize, postgres_type: &Type, ) -> Result, Error> { let path = match path { @@ -101,7 +130,7 @@ impl Bind { }, }; - let Some(param) = self.param_values.get(idx) else { + let Some(param) = self.param_values.get(value) else { return Ok(None); }; @@ -112,7 +141,6 @@ impl Bind { debug!( target: MAPPER, msg = "Fused JSON value selector", - ?idx, ?path, ?value ); @@ -120,74 +148,71 @@ impl Bind { Ok(Some(json_value_selector_plaintext(&path, value)?)) } - /// The one-plaintext-per-param pass. + /// Replaces the bound params with the output params of the rewritten + /// statement. /// - /// Params named in `json_value_selectors` are skipped: a fused value - /// selector has no plaintext of its own — its own bytes are only half the - /// needle — and decoding it here would fail, since the scalar it carries is - /// not a valid standalone operand for the column. - fn to_plaintext_positional( - &self, - param_columns: &[Option], - param_types: &[i32], - json_value_selectors: &HashMap, - ) -> Result>, Error> { - let plaintexts = self - .param_values - .iter() - .zip(param_columns.iter()) - .enumerate() - .map(|(idx, (param, col))| match col { - _ if json_value_selectors.contains_key(&idx) => Ok(None), - Some(col) => { - let bound_param_type = get_param_type(idx, param_types, col); - - debug!( - target: MAPPER, - col = ?col, bound_param_type = ?bound_param_type - ); - - // Convert param bytes into a Plaintext wrapping a Value - // If the param type is different, will convert the bound type to the correct Plaintext variant identified by the cast_type - let plaintext = bind_param_from_sql( - param, - &bound_param_type, - col.eql_term(), - col.cast_type(), - ) - .map_err(|_| MappingError::InvalidParameter(Box::new(col.to_owned())))?; - - Ok(plaintext) - } - None => Ok(None), - }) - .collect::, Error>>()?; - Ok(plaintexts) + /// When the plan is positional (the overwhelmingly common case) the params + /// are patched in place, leaving the client's framing — including its format + /// code encoding — exactly as sent. When the rewrite reshaped the params, + /// the list is rebuilt: each output param inherits the wire bytes and format + /// code of the input it was built around, and an explicit format code is + /// emitted per param since the counts no longer line up. + pub fn rewrite( + &mut self, + output_params: &[OutputParam], + encrypted: Vec>, + ) -> Result<(), Error> { + if output_params.len() == self.param_values.len() && params_are_positional(output_params) { + for (param, ct) in self.param_values.iter_mut().zip(encrypted.iter()) { + Self::apply_encrypted(param, ct.as_ref())?; + } + return Ok(()); + } + + let mut param_values = Vec::with_capacity(output_params.len()); + for (output, ct) in output_params.iter().zip(encrypted.iter()) { + let input = output.source.primary_input(); + let mut param = self.param_values.get(input).cloned().ok_or( + ProtocolError::MissingBoundParameter { + param: input + 1, + received: self.param_values.len(), + }, + )?; + + Self::apply_encrypted(&mut param, ct.as_ref())?; + param_values.push(param); + } + + self.param_format_codes = param_values.iter().map(|param| param.format_code).collect(); + self.num_param_format_codes = self.param_format_codes.len() as i16; + self.num_param_values = param_values.len() as i16; + self.param_values = param_values; + self.reshaped = true; + + Ok(()) } - pub fn rewrite(&mut self, encrypted: Vec>) -> Result<(), Error> { - for (idx, ct) in encrypted.iter().enumerate() { - if let Some(ct) = ct { - match ct { - // A JSON selector (`->`/`->>`/`jsonb_path_query`) is a bare - // tokenized-selector hash bound directly as `text`, NOT jsonb. - // Use the raw token: JSON-serializing it re-quotes the bare - // string (`""`), which never matches the stored per-entry - // `s`. It must also skip the jsonb version header a binary - // rewrite would prepend — the binary wire form of `text` is - // just its raw bytes, and a leading `0x01` corrupts the - // selector so `->` matches nothing. - EqlOutput::Query(EqlQueryPayload::Selector(s)) => { - self.param_values[idx].rewrite_text(s.clone().into_bytes()); - } - // convert json to bytes - _ => { - let bytes = serde_json::to_value(ct)?.to_string().into_bytes(); - self.param_values[idx].rewrite(&bytes); - } - } + fn apply_encrypted(param: &mut BindParam, ct: Option<&EqlOutput>) -> Result<(), Error> { + match ct { + // A JSON selector (`->`/`->>`/`jsonb_path_query`) is a bare + // tokenized-selector hash bound directly as `text`, NOT jsonb. + // Use the raw token: JSON-serializing it re-quotes the bare + // string (`""`), which never matches the stored per-entry + // `s`. It must also skip the jsonb version header a binary + // rewrite would prepend — the binary wire form of `text` is + // just its raw bytes, and a leading `0x01` corrupts the + // selector so `->` matches nothing. + Some(EqlOutput::Query(EqlQueryPayload::Selector(s))) => { + param.rewrite_text(s.clone().into_bytes()); } + // convert json to bytes + Some(ct) => { + let bytes = serde_json::to_value(ct)?.to_string().into_bytes(); + param.rewrite(&bytes); + } + None => {} } + Ok(()) } } @@ -369,6 +394,7 @@ impl TryFrom<&BytesMut> for Bind { param_values, num_result_column_format_codes, result_columns_format_codes, + reshaped: false, }) } } diff --git a/packages/cipherstash-proxy/src/postgresql/messages/mod.rs b/packages/cipherstash-proxy/src/postgresql/messages/mod.rs index f52c69d0c..b271c001c 100644 --- a/packages/cipherstash-proxy/src/postgresql/messages/mod.rs +++ b/packages/cipherstash-proxy/src/postgresql/messages/mod.rs @@ -24,6 +24,10 @@ pub use target::Target; pub const NULL: i32 = -1; +/// PostgreSQL's "unspecified type, infer it" param OID, used in `Parse` and +/// when a param's type is not known to the proxy. +pub const UNSPECIFIED_TYPE_OID: i32 = 0; + #[derive(Clone, Copy, Debug, PartialEq)] pub enum FrontendCode { Bind, diff --git a/packages/cipherstash-proxy/src/postgresql/messages/param_description.rs b/packages/cipherstash-proxy/src/postgresql/messages/param_description.rs index 56ba751e6..6a9b4c19b 100644 --- a/packages/cipherstash-proxy/src/postgresql/messages/param_description.rs +++ b/packages/cipherstash-proxy/src/postgresql/messages/param_description.rs @@ -46,6 +46,21 @@ impl ParamDescription { } } + /// Replaces the described params wholesale. + /// + /// PostgreSQL describes the params of the *rewritten* statement, but the + /// client must be told about the params it wrote — a rewrite that fuses two + /// params into one would otherwise describe too few, and the client would + /// bind the wrong number of values. + pub fn set_types(&mut self, types: Vec) { + debug!(target: MAPPER, ?types); + + if types != self.types { + self.types = types; + self.dirty = true; + } + } + pub fn requires_rewrite(&self) -> bool { self.dirty } diff --git a/packages/cipherstash-proxy/src/postgresql/messages/parse.rs b/packages/cipherstash-proxy/src/postgresql/messages/parse.rs index f89073ed4..e356b54bf 100644 --- a/packages/cipherstash-proxy/src/postgresql/messages/parse.rs +++ b/packages/cipherstash-proxy/src/postgresql/messages/parse.rs @@ -1,16 +1,13 @@ -use super::{FrontendCode, Name}; +use super::{FrontendCode, Name, UNSPECIFIED_TYPE_OID}; use crate::{ error::{Error, ProtocolError}, - postgresql::{context::column::Column, protocol::BytesMutReadString}, + postgresql::{context::statement::OutputParam, protocol::BytesMutReadString}, SIZE_I16, SIZE_I32, }; use bytes::{Buf, BufMut, BytesMut}; use postgres_types::Type; use std::{ffi::CString, io::Cursor}; -/// PostgreSQL's "unspecified type, infer it" param OID in a Parse message. -const UNSPECIFIED_TYPE_OID: i32 = 0; - #[derive(Debug, Clone)] pub struct Parse { pub code: char, @@ -26,52 +23,43 @@ impl Parse { self.dirty } + /// Rewrites the declared param types to describe the params of the + /// *rewritten* statement. /// /// EQL v3 encrypted columns are JSONB-backed domain types (e.g. - /// `eql_v3_text_search`). - /// - /// Using JSONB to avoid the complexity of loading each domain's OID — - /// PostgreSQL coerces JSONB to the domain type if it passes the CHECK - /// constraint. - /// - pub fn rewrite_param_types(&mut self, columns: &[Option]) { - for (idx, col) in columns.iter().enumerate() { - if self.param_types.get(idx).is_some() && col.is_some() { - self.param_types[idx] = Type::JSONB.oid() as i32; - self.dirty = true; - } - } - } - - /// - /// Declares a param type for every placeholder the rewrite left - /// unreferenced. - /// - /// A JSON field equality (`col -> $1 = $2`) fuses both operands into one - /// needle, so `$1` disappears from the rewritten SQL while staying in the - /// client's Bind. PostgreSQL allows an unreferenced param, but only if its - /// type is declared — otherwise it fails with "could not determine data type - /// of parameter $1". Clients that let the server infer types (the common - /// case) send no types at all, so the array is grown here. + /// `eql_v3_text_search`). JSONB is declared rather than the domain itself to + /// avoid loading each domain's OID — PostgreSQL coerces JSONB to the domain + /// if it passes the CHECK constraint. /// - /// Slots that are still referenced are left as OID 0, which PostgreSQL reads - /// as "unspecified, infer it" — so this never overrides a type the query - /// itself determines. The dropped selectors are bound as the encrypted - /// selector *text*, hence TEXT. + /// The client declares types for the params it wrote; the rewrite may have + /// dropped or fused some of those, so each declaration is carried across to + /// the output param that consumes it. An output param that carries an + /// encrypted value is declared JSONB regardless — that is the wire type of + /// every EQL payload, whatever the client thought it was binding. /// - pub fn declare_unreferenced_param_types(&mut self, indexes: &[usize]) { - let Some(required_len) = indexes.iter().max().map(|idx| idx + 1) else { + /// A client that declares no types at all (the common case — it lets the + /// server infer them) is left alone: every output param is referenced by the + /// rewritten SQL, so PostgreSQL can always infer them. + pub fn rewrite_param_types(&mut self, output_params: &[OutputParam]) { + if self.param_types.is_empty() { return; - }; - - if self.param_types.len() < required_len { - self.param_types.resize(required_len, UNSPECIFIED_TYPE_OID); - self.num_params = self.param_types.len() as i16; - self.dirty = true; } - for idx in indexes { - self.param_types[*idx] = Type::TEXT.oid() as i32; + let param_types = output_params + .iter() + .map(|output| match &output.column { + Some(_) => Type::JSONB.oid() as i32, + None => self + .param_types + .get(output.source.primary_input()) + .copied() + .unwrap_or(UNSPECIFIED_TYPE_OID), + }) + .collect::>(); + + if param_types != self.param_types { + self.num_params = param_types.len() as i16; + self.param_types = param_types; self.dirty = true; } } @@ -156,7 +144,11 @@ mod tests { use crate::{ config::LogConfig, log, - postgresql::{messages::parse::Parse, Column}, + postgresql::{ + context::statement::{OutputParam, OutputParamSource}, + messages::parse::Parse, + Column, + }, Identifier, }; use bytes::BytesMut; @@ -195,9 +187,52 @@ mod tests { let config = ColumnConfig::build("column".to_string()).casts_as(ColumnType::SmallInt); let column = Column::new(identifier, config, None, eql_mapper::EqlTermVariant::Full); - let columns = vec![None, Some(column)]; + let output_params = vec![ + OutputParam { + column: None, + source: OutputParamSource::Input(0), + }, + OutputParam { + column: Some(column), + source: OutputParamSource::Input(1), + }, + ]; + + parse.rewrite_param_types(&output_params); + assert!(parse.requires_rewrite()); + assert_eq!( + parse.param_types, + vec![ + postgres_types::Type::INT2.oid() as i32, + postgres_types::Type::JSONB.oid() as i32 + ] + ); + } + + /// A rewrite that fuses two params into one must leave the client's + /// declaration for the surviving param, not the one it happened to sit at. + #[test] + pub fn test_parse_rewrite_param_types_after_fusion() { + log::init(LogConfig::default()); + let bytes = to_message( + b"P\0\0\0J\0INSERT INTO encrypted (id, encrypted_int2) VALUES ($1, $2)\0\0\x02\0\0\0\x15\0\0\0\x15" + ); + + let mut parse = Parse::try_from(&bytes).unwrap(); + + // Two input params collapse to a single native output param sourced + // from input 1. + let output_params = vec![OutputParam { + column: None, + source: OutputParamSource::Input(1), + }]; - parse.rewrite_param_types(&columns); + parse.rewrite_param_types(&output_params); assert!(parse.requires_rewrite()); + assert_eq!(parse.num_params, 1); + assert_eq!( + parse.param_types, + vec![postgres_types::Type::INT2.oid() as i32] + ); } } diff --git a/packages/eql-mapper/src/lib.rs b/packages/eql-mapper/src/lib.rs index 0f4d176ea..eae3c93eb 100644 --- a/packages/eql-mapper/src/lib.rs +++ b/packages/eql-mapper/src/lib.rs @@ -9,6 +9,8 @@ mod iterator_ext; mod json_value_selector; mod model; mod param; +mod param_plan; +mod renumber_params; mod scope_tracker; mod transformation_rules; mod type_checked_statement; @@ -21,6 +23,7 @@ pub use eql_mapper::*; pub use json_value_selector::*; pub use model::*; pub use param::*; +pub use param_plan::*; pub use type_checked_statement::*; pub use unifier::{ Array, AssociatedType, DomainIdentity, EqlTerm, EqlTermVariant, EqlTrait, EqlTraits, EqlValue, @@ -29,6 +32,7 @@ pub use unifier::{ pub(crate) use dep::*; pub(crate) use inference::*; +pub(crate) use renumber_params::*; pub(crate) use scope_tracker::*; pub(crate) use transformation_rules::*; @@ -41,7 +45,7 @@ mod test { EqlTerm, EqlTrait, EqlTraits, EqlValue, InstantiateType, NativeValue, Projection, ProjectionColumn, Type, Value, }, - JsonSelectorSource, Param, Schema, TableColumn, TableResolver, + JsonSelectorSource, OutputParamSource, Param, Schema, TableColumn, TableResolver, }; use eql_mapper_macros::concrete_ty; use pretty_assertions::assert_eq; @@ -2368,26 +2372,30 @@ mod test { } /// `col -> $1 = $2` fuses both operands into ONE containment needle: the - /// field access is discarded, the value operand becomes the needle, and the - /// selector placeholder is left declared-but-unreferenced (which keeps param - /// numbering identical between client and server). + /// field access is discarded and the two input params become a single + /// output param, renumbered `$1`. #[test] fn json_field_eq_params_rewrites_to_containment() { let statement = parse("SELECT id FROM patients WHERE notes -> $1 = $2"); let typed = type_check(json_eq_schema(), &statement).unwrap(); + let transformed = typed.transform(HashMap::new()).unwrap(); assert_eq!( - typed.transform(HashMap::new()).unwrap().to_string(), - "SELECT id FROM patients WHERE eql_v3.jsonb_contains(notes, $2::JSONB::eql_v3.query_json)" + transformed.to_string(), + "SELECT id FROM patients WHERE eql_v3.jsonb_contains(notes, $1::JSONB::eql_v3.query_json)" ); - // The value operand ($2) carries the needle; its path comes from $1. + // Two input params, one output param, derived from both. + assert_eq!(transformed.params.len(), 1); assert_eq!( - typed.json_value_selectors.for_param(Param(2)), - Some(&JsonSelectorSource::Param(Param(1))) + transformed.params.outputs()[0].source, + OutputParamSource::JsonValueSelector { + path: JsonSelectorSource::Param(Param(1)), + value: Param(2), + } ); - assert_eq!(typed.json_value_selectors.for_param(Param(1)), None); + assert!(!transformed.params.is_identity()); } /// The `->>` and `jsonb_path_query_first` spellings are the same access and @@ -2400,12 +2408,61 @@ mod test { assert_eq!( typed.transform(HashMap::new()).unwrap().to_string(), - "SELECT id FROM patients WHERE eql_v3.jsonb_contains(notes, $2::JSONB::eql_v3.query_json)", + "SELECT id FROM patients WHERE eql_v3.jsonb_contains(notes, $1::JSONB::eql_v3.query_json)", "unexpected rewrite for `{access}`" ); } } + /// Params around a fusion are renumbered to close the gap it leaves, and the + /// plan records where each surviving output param came from. + #[test] + fn json_field_eq_renumbers_surrounding_params() { + let statement = + parse("SELECT id FROM patients WHERE id = $1 AND notes -> $2 = $3 AND id <> $4"); + + let typed = type_check(json_eq_schema(), &statement).unwrap(); + let transformed = typed.transform(HashMap::new()).unwrap(); + + assert_eq!( + transformed.to_string(), + "SELECT id FROM patients WHERE id = $1 AND eql_v3.jsonb_contains(notes, $2::JSONB::eql_v3.query_json) AND id <> $3" + ); + + let outputs = transformed.params.outputs(); + assert_eq!(outputs.len(), 3); + assert_eq!(outputs[0].source, OutputParamSource::Input(Param(1))); + assert_eq!( + outputs[1].source, + OutputParamSource::JsonValueSelector { + path: JsonSelectorSource::Param(Param(2)), + value: Param(3), + } + ); + assert_eq!(outputs[2].source, OutputParamSource::Input(Param(4))); + } + + /// A statement whose params the rewrite leaves alone reports an identity + /// plan, so the proxy can keep binding by position. + #[test] + fn ordinary_params_produce_an_identity_plan() { + let schema = resolver(schema! { + tables: { + patients: { + id, + name (EQL: Eq), + } + } + }); + + let statement = parse("SELECT id FROM patients WHERE id = $1 AND name = $2"); + let typed = type_check(schema, &statement).unwrap(); + let transformed = typed.transform(HashMap::new()).unwrap(); + + assert_eq!(transformed.params.len(), 2); + assert!(transformed.params.is_identity()); + } + /// Both halves as literals: the selector text is captured inline at /// type-check time, and the selector literal vanishes from the output. #[test] @@ -2457,7 +2514,7 @@ mod test { assert_eq!( typed.transform(HashMap::new()).unwrap().to_string(), - "SELECT id FROM patients WHERE NOT (eql_v3.jsonb_contains(notes, $2::JSONB::eql_v3.query_json))" + "SELECT id FROM patients WHERE NOT (eql_v3.jsonb_contains(notes, $1::JSONB::eql_v3.query_json))" ); } diff --git a/packages/eql-mapper/src/param_plan.rs b/packages/eql-mapper/src/param_plan.rs new file mode 100644 index 000000000..877d58a4d --- /dev/null +++ b/packages/eql-mapper/src/param_plan.rs @@ -0,0 +1,120 @@ +//! The correspondence between the params of the input statement and the params +//! of the rewritten one. +//! +//! There is no 1:1 guarantee. An output param may be **derived from more than +//! one input param** — encrypted JSON equality (`col -> $1 = $2`) fuses a path +//! and a value into a single value-selector needle, so two input params become +//! one output param. Nothing in the pipeline may assume that param `$n` on the +//! wire to PostgreSQL is param `$n` as the client wrote it. +//! +//! A [`ParamPlan`] is the explicit statement of that correspondence: one +//! [`OutputParam`] per placeholder in the rewritten SQL, each naming the input +//! params its value is built from. The proxy binds against the plan — it +//! describes the *input* params to the client, and sends the *output* params to +//! PostgreSQL. + +use std::collections::HashSet; + +use crate::unifier::Value; +use crate::{EqlMapperError, JsonSelectorSource, Param}; + +/// How the value bound to one output param is derived from the input params. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum OutputParamSource { + /// Forwarded from a single input param, unchanged in meaning (though still + /// encrypted if the param is EQL-typed). + Input(Param), + + /// Fused from two operands into one encrypted value-selector needle: the + /// JSON path and the value it must equal. The path may itself be a param or + /// a literal in the SQL; the value is always the param carrying this output. + /// + /// See [`crate::JsonValueSelectors`]. + JsonValueSelector { + path: JsonSelectorSource, + value: Param, + }, +} + +impl OutputParamSource { + /// Every input param this output param consumes. + pub fn inputs(&self) -> Vec { + match self { + OutputParamSource::Input(param) => vec![*param], + OutputParamSource::JsonValueSelector { path, value } => match path { + JsonSelectorSource::Param(path_param) => vec![*path_param, *value], + JsonSelectorSource::Literal(_) => vec![*value], + }, + } + } +} + +/// One placeholder of the rewritten statement. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct OutputParam { + /// Its 1-based position in the rewritten SQL. + pub param: Param, + + /// The type of the value to bind — what the proxy must encrypt it as. + pub value: Value, + + /// The input params it is built from. + pub source: OutputParamSource, +} + +/// The params of the rewritten statement, in the order PostgreSQL will see them. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct ParamPlan { + outputs: Vec, +} + +impl ParamPlan { + pub(crate) fn new(outputs: Vec) -> Self { + Self { outputs } + } + + pub fn outputs(&self) -> &[OutputParam] { + &self.outputs + } + + pub fn len(&self) -> usize { + self.outputs.len() + } + + pub fn is_empty(&self) -> bool { + self.outputs.is_empty() + } + + /// `true` when every input param maps to the output param of the same + /// position — the ordinary case, where the rewrite changed no param. + pub fn is_identity(&self) -> bool { + self.outputs.iter().enumerate().all(|(idx, output)| { + matches!(output.source, OutputParamSource::Input(input) + if input.0 as usize == idx + 1) + }) + } + + /// Checks that the plan consumes every input param. + /// + /// The invariant the rewrite must uphold is **coverage, not exactly-once**: + /// an input may feed several output params (if the rewrite duplicated it), + /// and several inputs may feed one output (fusion), but an input that feeds + /// nothing has been silently dropped — the client would bind a value that + /// never reaches PostgreSQL and never contributes to a needle, which is a + /// bug in a rewrite rule rather than a valid query. + pub(crate) fn check_covers(&self, inputs: &[Param]) -> Result<(), EqlMapperError> { + let consumed: HashSet = self + .outputs + .iter() + .flat_map(|output| output.source.inputs()) + .collect(); + + if let Some(orphan) = inputs.iter().find(|param| !consumed.contains(param)) { + return Err(EqlMapperError::InternalError(format!( + "param {orphan} of the input statement is not consumed by the rewritten statement" + ))); + } + + Ok(()) + } +} diff --git a/packages/eql-mapper/src/renumber_params.rs b/packages/eql-mapper/src/renumber_params.rs new file mode 100644 index 000000000..f8d15e2a2 --- /dev/null +++ b/packages/eql-mapper/src/renumber_params.rs @@ -0,0 +1,60 @@ +//! Renumbering of the rewritten statement's placeholders. +//! +//! A rewrite rule may drop a placeholder (encrypted JSON equality folds the +//! path operand into the value's needle) or duplicate one. Either way the +//! surviving placeholders no longer read `$1..$n` in order, and PostgreSQL +//! requires a statement's params to be exactly `$1..$m` — an unreferenced `$n` +//! is only accepted when its type is declared, and a gap is a statement whose +//! param count no longer matches what the client bound. +//! +//! So after the rewrite rules have run, this pass walks the new statement in +//! SQL order and renumbers the placeholders `$1..$m`, recording which input +//! param each output position came from. That record is the raw material for a +//! [`crate::ParamPlan`]. +//! +//! This runs as a **separate pass, after** the rewrite rules, which is what +//! lets [`crate::FailOnPlaceholderChange`] keep enforcing that no rule replaces +//! a placeholder with a literal during the rewrite itself. + +use sqltk::parser::ast; +use sqltk::{NodePath, Transform, Visitable}; + +use crate::{EqlMapperError, Param}; + +/// Assigns `$1..$m` to the placeholders of a rewritten statement in SQL order. +#[derive(Debug, Default)] +pub(crate) struct RenumberParams { + /// The input param each output position was renumbered from, in output + /// order. Index `i` holds the source of `$(i + 1)`. + sources: Vec, +} + +impl RenumberParams { + pub(crate) fn new() -> Self { + Self::default() + } + + pub(crate) fn into_sources(self) -> Vec { + self.sources + } +} + +impl<'ast> Transform<'ast> for RenumberParams { + type Error = EqlMapperError; + + fn transform( + &mut self, + _node_path: &NodePath<'ast>, + mut target_node: N, + ) -> Result { + // Transformation is depth-first and visits siblings in order, so + // placeholders arrive in the order they appear in the rendered SQL. + if let Some(ast::Value::Placeholder(name)) = target_node.downcast_mut::() { + let source = Param::try_from(&*name)?; + self.sources.push(source); + *name = format!("${}", self.sources.len()); + } + + Ok(target_node) + } +} diff --git a/packages/eql-mapper/src/type_checked_statement.rs b/packages/eql-mapper/src/type_checked_statement.rs index d462f4675..ef12d807e 100644 --- a/packages/eql-mapper/src/type_checked_statement.rs +++ b/packages/eql-mapper/src/type_checked_statement.rs @@ -6,13 +6,40 @@ use sqltk::{AsNodeKey, NodeKey, Transformable}; use crate::unifier::EqlTerm; use crate::{ CastLiteralsAsEncrypted, CastParamsAsEncrypted, DryRunnable, EqlMapperError, - FailOnPlaceholderChange, JsonValueSelectors, Param, PreserveEffectiveAliases, - RewriteContainmentOps, RewriteEqlComparisonOps, RewriteEqlMatchOps, RewriteJsonValueSelectorEq, - RewriteStandardSqlFnsOnEqlTypes, TransformationRule, + FailOnPlaceholderChange, JsonValueSelectors, OutputParam, OutputParamSource, Param, ParamPlan, + PreserveEffectiveAliases, RenumberParams, RewriteContainmentOps, RewriteEqlComparisonOps, + RewriteEqlMatchOps, RewriteJsonValueSelectorEq, RewriteStandardSqlFnsOnEqlTypes, + TransformationRule, }; use crate::unifier::{Projection, Type, Value}; +/// The result of [`TypeCheckedStatement::transform`]: the rewritten statement +/// and the correspondence between its params and the input's. +/// +/// Derefs to the [`Statement`] so a caller that only wants the SQL can treat it +/// as one; a caller that binds params must consult [`Self::params`], because +/// the two param lists are not guaranteed to correspond by position. +#[derive(Debug)] +pub struct TransformedStatement { + pub statement: Statement, + pub params: ParamPlan, +} + +impl std::ops::Deref for TransformedStatement { + type Target = Statement; + + fn deref(&self) -> &Self::Target { + &self.statement + } +} + +impl std::fmt::Display for TransformedStatement { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.statement.fmt(f) + } +} + /// A `TypeCheckedStatement` is returned from a successful call to [`crate::type_check`]. #[derive(Debug)] pub struct TypeCheckedStatement<'ast> { @@ -32,11 +59,11 @@ pub struct TypeCheckedStatement<'ast> { /// [`EqlTerm::JsonValueSelector`], where the path half of its needle comes /// from. /// - /// This is the one place the mapper's output is **not** 1:1 with the input - /// SQL. `col -> sel = value` consumes two operands and emits one encrypted - /// needle; the path operand is dropped from the rewritten statement (its - /// placeholder stays declared but unreferenced, so param numbering is - /// untouched). The proxy consults this to compose the needle. + /// This is why the mapper's output is **not** 1:1 with the input SQL: + /// `col -> sel = value` consumes two operands and emits one encrypted + /// needle. The path operand is dropped from the rewritten statement, and + /// the remaining params are renumbered — see [`ParamPlan`], which is what + /// the proxy binds against. pub json_value_selectors: JsonValueSelectors<'ast>, /// A [`HashMap`] of AST node (using [`NodeKey`] as the key) to [`Type`]. The map contains a `Type` for every node @@ -99,14 +126,72 @@ impl<'ast> TypeCheckedStatement<'ast> { /// Transforms the SQL statement by replacing all plaintext literals with EQL equivalents /// and inserting EQL helper functions where necessary. + /// + /// Returns the rewritten statement together with its [`ParamPlan`] — the + /// params of the rewritten statement are **not** guaranteed to correspond + /// 1:1 with the input's, so the caller must bind through the plan rather + /// than by position. pub fn transform( &self, encrypted_literals: HashMap, sqltk::parser::ast::Value>, - ) -> Result { + ) -> Result { self.check_all_encrypted_literals_provided(&encrypted_literals)?; let mut transformer = self.make_transformer(encrypted_literals); transformer.set_real_run_mode(); - self.statement.apply_transform(&mut transformer) + let rewritten = self.statement.apply_transform(&mut transformer)?; + + // Renumber in a second pass, so the rules above are free to drop or + // duplicate placeholders without having to maintain `$n` themselves — + // and so `FailOnPlaceholderChange` still governs the rules' own edits. + let mut renumber = RenumberParams::new(); + let statement = rewritten.apply_transform(&mut renumber)?; + let params = self.param_plan(renumber.into_sources())?; + + Ok(TransformedStatement { statement, params }) + } + + /// Builds the [`ParamPlan`] from the input param each output placeholder was + /// renumbered from. + /// + /// An output param whose input is a fused JSON value selector carries both + /// operands; every other output forwards its input alone. + fn param_plan(&self, sources: Vec) -> Result { + let outputs = sources + .into_iter() + .enumerate() + .map(|(idx, input)| { + let value = self + .params + .iter() + .find_map(|(param, value)| (*param == input).then(|| value.clone())) + .ok_or_else(|| { + EqlMapperError::InternalError(format!( + "rewritten statement refers to param {input}, which the input statement does not declare" + )) + })?; + + let source = match self.json_value_selectors.for_param(input) { + Some(path) => OutputParamSource::JsonValueSelector { + path: path.clone(), + value: input, + }, + None => OutputParamSource::Input(input), + }; + + Ok(OutputParam { + param: Param((idx + 1) as u16), + value, + source, + }) + }) + .collect::, EqlMapperError>>()?; + + let plan = ParamPlan::new(outputs); + + let inputs: Vec = self.params.iter().map(|(param, _)| *param).collect(); + plan.check_covers(&inputs)?; + + Ok(plan) } pub fn literal_values(&self) -> &Vec<(EqlTerm, &'ast sqltk::parser::ast::Value)> { From 51435b03f1564713f8eaf74bc03d2fbd42c9c31e Mon Sep 17 00:00:00 2001 From: James Sadler Date: Sun, 26 Jul 2026 10:58:20 +1000 Subject: [PATCH 05/10] refactor(mapper): each rewrite rule casts its own operands MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Removes `CastLiteralsAsEncrypted` and `CastParamsAsEncrypted`. Both were context-blind: they fired on any encrypted literal or placeholder and then guessed which domain to cast it to by walking up the ancestor chain (`is_query_operand`) looking for a comparison or a LIKE. The rule that actually knew the context — the one rewriting the predicate — was a different rule entirely, running later. Now the rule that owns a construct casts the operands of that construct, where the context is a fact rather than an inference: - `RewriteEqlComparisonOps`, `RewriteEqlMatchOps` — query operands, cast to the term-only `eql_v3.query_*` twin. - `RewriteJsonValueSelectorEq` — the fused needle, `eql_v3.query_json`. - `RewriteContainmentOps` — a `@>`/`<@` needle is a whole document, so it casts to the column domain; a `->`/`->>` selector takes no cast. - `CastFullPayloadOperands` (new) — the contexts with no rewrite of their own: INSERT values, UPDATE assignments, and directly-written `eql_v3.jsonb_contains(col, $1)` calls, which clients on platforms without operator support write themselves and which need the column-domain cast for the GIN index over `eql_v3.jsonb_array(col)`. It fires on the enclosing `Values`/`Assignment`/`Function` node, not on the literal, so the context is again structural. Substitution is separated from casting: `SubstituteEncryptedLiterals` puts each ciphertext in place wherever it appears and casts nothing, since where a literal sits decides its domain but not whether it needs replacing. It keeps the postcondition that every encrypted literal was consumed. The two domain choices are now named for the distinction they encode — `query_operand_domain` (just the terms a predicate needs) and `full_payload_domain` (ciphertext plus every indexed term) — and `is_query_operand`, the ancestor walk, is gone. Every rule matches against the ORIGINAL node via `node_path`, so a construct already rewritten by an earlier rule cannot be cast twice. Mapper 104 unit tests and proxy 117 unchanged. Integration: no regressions — the only tests that differ from the previous run are in the flaky ORE/OPE ordering family (net +1 passing), and the one that regressed both appears in the base binary's failure list and flips between ok and failed across identical isolated runs. Stable-Commit-Id: q-3y30vhx105f1ms7h36tanr8k7m --- .../cast_full_payload_operands.rs | 196 ++++++++++++++++++ .../cast_literals_as_encrypted.rs | 103 --------- .../cast_params_as_encrypted.rs | 103 --------- .../src/transformation_rules/helpers.rs | 148 +++++++------ .../src/transformation_rules/mod.rs | 8 +- .../rewrite_containment_ops.rs | 27 ++- .../rewrite_eql_comparison_ops.rs | 25 ++- .../rewrite_eql_match_ops.rs | 42 +++- .../rewrite_json_value_selector_eq.rs | 21 +- .../substitute_encrypted_literals.rs | 74 +++++++ .../eql-mapper/src/type_checked_statement.rs | 16 +- 11 files changed, 458 insertions(+), 305 deletions(-) create mode 100644 packages/eql-mapper/src/transformation_rules/cast_full_payload_operands.rs delete mode 100644 packages/eql-mapper/src/transformation_rules/cast_literals_as_encrypted.rs delete mode 100644 packages/eql-mapper/src/transformation_rules/cast_params_as_encrypted.rs create mode 100644 packages/eql-mapper/src/transformation_rules/substitute_encrypted_literals.rs diff --git a/packages/eql-mapper/src/transformation_rules/cast_full_payload_operands.rs b/packages/eql-mapper/src/transformation_rules/cast_full_payload_operands.rs new file mode 100644 index 000000000..0531d0779 --- /dev/null +++ b/packages/eql-mapper/src/transformation_rules/cast_full_payload_operands.rs @@ -0,0 +1,196 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use sqltk::parser::ast::{ + Assignment, Expr, Function, FunctionArg, FunctionArgExpr, FunctionArguments, Values, +}; +use sqltk::{NodeKey, NodePath, Visitable}; + +use crate::unifier::{Type, Value}; +use crate::EqlMapperError; + +use super::helpers::{cast_encrypted_operand, full_payload_domain}; +use super::TransformationRule; + +/// Casts encrypted values that must carry the column's **whole payload** — the +/// ciphertext plus every search term the column indexes — to the column domain, +/// rather than to a term-only `eql_v3.query_*` twin. +/// +/// Three contexts need it, and none of them is a predicate whose own rewrite +/// rule could own the cast: +/// +/// - `INSERT INTO t (col) VALUES ($1)` — the value is stored. +/// - `UPDATE t SET col = 'x'` — likewise. +/// - `eql_v3.jsonb_contains(col, $1)` and friends — a containment needle is a +/// whole document, and the cast is what lets PostgreSQL use the GIN index over +/// `eql_v3.jsonb_array(col)`. Clients on platforms without operator support +/// (Supabase, PostgREST) write these function forms directly, so they arrive +/// already spelled `eql_v3.*` with no operator for a rewrite rule to catch. +/// +/// The rule fires on the enclosing `Values`, `Assignment` and `Function` nodes +/// rather than on the value expressions, so "this operand carries a full +/// payload" is a fact about the construct that owns it, not a guess made by +/// walking up the tree from a literal. +/// +/// A JSON selector argument is left uncast — [`full_payload_domain`] returns +/// `None` for it, because `eql_v3.jsonb_path_query(json, text)` takes the bare +/// encrypted selector text. +#[derive(Debug)] +pub struct CastFullPayloadOperands<'ast> { + node_types: Arc, Type>>, +} + +impl<'ast> CastFullPayloadOperands<'ast> { + pub fn new(node_types: Arc, Type>>) -> Self { + Self { node_types } + } + + /// Whether `expr` is an encrypted literal or placeholder that this rule + /// would cast. Shared by `apply` and `would_edit` so the dry run agrees with + /// the real run. + fn needs_cast(&self, expr: &'ast Expr) -> bool { + matches!(expr, Expr::Value(_)) + && matches!( + self.node_types.get(&NodeKey::new(expr)), + Some(Type::Value(Value::Eql(eql_term))) if full_payload_domain(eql_term).is_some() + ) + } + + /// The argument expressions of a function call, in order. + fn args(function: &Function) -> impl Iterator { + let args = match &function.args { + FunctionArguments::List(list) => Some(list.args.iter()), + _ => None, + }; + + args.into_iter().flatten().filter_map(|arg| match arg { + FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) + | FunctionArg::Named { + arg: FunctionArgExpr::Expr(expr), + .. + } + | FunctionArg::ExprNamed { + arg: FunctionArgExpr::Expr(expr), + .. + } => Some(expr), + _ => None, + }) + } + + fn args_mut(function: &mut Function) -> impl Iterator { + let args = match &mut function.args { + FunctionArguments::List(list) => Some(list.args.iter_mut()), + _ => None, + }; + + args.into_iter().flatten().filter_map(|arg| match arg { + FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) + | FunctionArg::Named { + arg: FunctionArgExpr::Expr(expr), + .. + } + | FunctionArg::ExprNamed { + arg: FunctionArgExpr::Expr(expr), + .. + } => Some(expr), + _ => None, + }) + } + + /// Whether `function` is an `eql_v3.*` call — the only functions whose + /// encrypted arguments this rule owns. + fn is_eql_v3_function(function: &Function) -> bool { + function + .name + .0 + .first() + .and_then(|part| part.as_ident()) + .is_some_and(|ident| ident.value.eq_ignore_ascii_case("eql_v3")) + } +} + +impl<'ast> TransformationRule<'ast> for CastFullPayloadOperands<'ast> { + fn apply( + &mut self, + node_path: &NodePath<'ast>, + target_node: &mut N, + ) -> Result { + if let Some((original,)) = node_path.last_1_as::() { + let Some(target) = target_node.downcast_mut::() else { + return Ok(false); + }; + + let mut edited = false; + for (original_row, target_row) in original.rows.iter().zip(target.rows.iter_mut()) { + for (original_expr, target_expr) in original_row.iter().zip(target_row.iter_mut()) { + edited |= cast_encrypted_operand( + &self.node_types, + original_expr, + target_expr, + full_payload_domain, + ); + } + } + + return Ok(edited); + } + + if let Some((original,)) = node_path.last_1_as::() { + let Some(target) = target_node.downcast_mut::() else { + return Ok(false); + }; + + return Ok(cast_encrypted_operand( + &self.node_types, + &original.value, + &mut target.value, + full_payload_domain, + )); + } + + if let Some((original,)) = node_path.last_1_as::() { + if !Self::is_eql_v3_function(original) { + return Ok(false); + } + + let Some(target) = target_node.downcast_mut::() else { + return Ok(false); + }; + + let mut edited = false; + for (original_arg, target_arg) in Self::args(original).zip(Self::args_mut(target)) { + edited |= cast_encrypted_operand( + &self.node_types, + original_arg, + target_arg, + full_payload_domain, + ); + } + + return Ok(edited); + } + + Ok(false) + } + + fn would_edit(&mut self, node_path: &NodePath<'ast>, _target_node: &N) -> bool { + if let Some((original,)) = node_path.last_1_as::() { + return original + .rows + .iter() + .flat_map(|row| row.iter()) + .any(|expr| self.needs_cast(expr)); + } + + if let Some((original,)) = node_path.last_1_as::() { + return self.needs_cast(&original.value); + } + + if let Some((original,)) = node_path.last_1_as::() { + return Self::is_eql_v3_function(original) + && Self::args(original).any(|expr| self.needs_cast(expr)); + } + + false + } +} diff --git a/packages/eql-mapper/src/transformation_rules/cast_literals_as_encrypted.rs b/packages/eql-mapper/src/transformation_rules/cast_literals_as_encrypted.rs deleted file mode 100644 index 53e68032f..000000000 --- a/packages/eql-mapper/src/transformation_rules/cast_literals_as_encrypted.rs +++ /dev/null @@ -1,103 +0,0 @@ -use std::{any::type_name, collections::HashMap, sync::Arc}; - -use sqltk::parser::ast::{Expr, Value, ValueWithSpan}; -use sqltk::parser::tokenizer::Span; -use sqltk::{NodeKey, NodePath, Visitable}; - -use crate::unifier::{EqlTerm, Type, Value as UnifierValue}; -use crate::EqlMapperError; - -use super::helpers::{cast_to_v3_domain, json_query_operand_cast_target, v3_cast_target}; -use super::TransformationRule; - -#[derive(Debug)] -pub struct CastLiteralsAsEncrypted<'ast> { - encrypted_literals: HashMap, Value>, - node_types: Arc, Type>>, -} - -impl<'ast> CastLiteralsAsEncrypted<'ast> { - pub fn new( - encrypted_literals: HashMap, Value>, - node_types: Arc, Type>>, - ) -> Self { - Self { - encrypted_literals, - node_types, - } - } -} - -impl<'ast> TransformationRule<'ast> for CastLiteralsAsEncrypted<'ast> { - fn apply( - &mut self, - node_path: &NodePath<'ast>, - target_node: &mut N, - ) -> Result { - if self.would_edit(node_path, target_node) { - if let Some((original @ Expr::Value(ValueWithSpan { value, .. }),)) = - node_path.last_1_as::() - { - if let Some(replacement) = self.encrypted_literals.remove(&NodeKey::new(value)) { - // The literal's domain identity determines the cast target; its - // role (query operand vs stored value) determines which v3 - // domain (query twin vs column domain). - let Some(Type::Value(UnifierValue::Eql(eql_term))) = - self.node_types.get(&NodeKey::new(original)) - else { - return Err(EqlMapperError::Transform(format!( - "{}: encrypted literal has no EQL type", - type_name::() - ))); - }; - - let target_node = target_node.downcast_mut::().unwrap(); - *target_node = match eql_term { - // A JSON selector — the RHS of `->`/`->>`, or the path - // argument of `jsonb_path_query` — is passed to the EQL v3 - // function as the encrypted-selector *text* - // (`eql_v3."->"(json, text)`, `jsonb_path_query(json, - // text)`), not a jsonb-domain payload. The encrypt pipeline - // produces that selector text (a SteVec `QueryOp::SteVecSelector` - // token), so it is emitted verbatim, uncast. - EqlTerm::JsonAccessor(_) | EqlTerm::JsonPath(_) => { - Expr::Value(ValueWithSpan { - value: replacement, - span: Span::empty(), - }) - } - _ => { - let (schema, domain) = json_query_operand_cast_target(eql_term) - .unwrap_or_else(|| { - let identity = eql_term.eql_value().domain_identity().clone(); - v3_cast_target(node_path, &identity) - }); - cast_to_v3_domain(replacement, &schema, &domain) - } - }; - return Ok(true); - } - } - } - - Ok(false) - } - - fn would_edit(&mut self, node_path: &NodePath<'ast>, _target_node: &N) -> bool { - if let Some((Expr::Value(ValueWithSpan { value, .. }),)) = node_path.last_1_as::() { - return self.encrypted_literals.contains_key(&NodeKey::new(value)); - } - false - } - - fn check_postcondition(&self) -> Result<(), EqlMapperError> { - if self.encrypted_literals.is_empty() { - Ok(()) - } else { - Err(EqlMapperError::Transform(format!( - "Postcondition failed in {}: unused encrypted literals", - type_name::() - ))) - } - } -} diff --git a/packages/eql-mapper/src/transformation_rules/cast_params_as_encrypted.rs b/packages/eql-mapper/src/transformation_rules/cast_params_as_encrypted.rs deleted file mode 100644 index 56278b690..000000000 --- a/packages/eql-mapper/src/transformation_rules/cast_params_as_encrypted.rs +++ /dev/null @@ -1,103 +0,0 @@ -use super::helpers::{cast_to_v3_domain, json_query_operand_cast_target, v3_cast_target}; -use super::TransformationRule; -use crate::unifier::{EqlTerm, Type, Value as UnifierValue}; -use crate::EqlMapperError; -use sqltk::parser::ast::{Expr, Value, ValueWithSpan}; -use sqltk::parser::tokenizer::Span; -use sqltk::{NodeKey, NodePath, Visitable}; -use std::collections::HashMap; -use std::sync::Arc; - -#[derive(Debug)] -pub struct CastParamsAsEncrypted<'ast> { - node_types: Arc, Type>>, -} - -impl<'ast> CastParamsAsEncrypted<'ast> { - pub fn new(node_types: Arc, Type>>) -> Self { - Self { node_types } - } -} - -impl<'ast> TransformationRule<'ast> for CastParamsAsEncrypted<'ast> { - fn apply( - &mut self, - node_path: &NodePath<'ast>, - target_node: &mut N, - ) -> Result { - if self.would_edit(node_path, target_node) { - // Resolve the operand's v3 cast target from its domain identity and - // its role (query operand → query twin; stored value → column domain). - let Some((original,)) = node_path.last_1_as::() else { - return Ok(false); - }; - let Some(Type::Value(UnifierValue::Eql(eql_term))) = - self.node_types.get(&NodeKey::new(original)) - else { - return Ok(false); - }; - - // A JSON selector operand (RHS of `->`/`->>`, or the path argument of - // `jsonb_path_query`) is passed to the EQL v3 function as the encrypted - // selector *text* — `eql_v3."->"(json, text)`, `jsonb_path_query(json, - // text)` — not a jsonb query-domain payload. The proxy encrypts these - // params as SteVec selectors, so the placeholder is left bare (its type - // is inferred from the function signature) rather than cast to the - // `eql_v3.query_*` twin. Mirrors the JsonAccessor arm of - // `CastLiteralsAsEncrypted`. - if matches!(eql_term, EqlTerm::JsonAccessor(_) | EqlTerm::JsonPath(_)) { - return Ok(false); - } - - let (schema, domain) = json_query_operand_cast_target(eql_term).unwrap_or_else(|| { - let identity = eql_term.eql_value().domain_identity().clone(); - v3_cast_target(node_path, &identity) - }); - - if let Some( - expr @ Expr::Value(ValueWithSpan { - value: Value::Placeholder(_), - .. - }), - ) = target_node.downcast_mut() - { - let to_wrap = std::mem::replace( - expr, - Expr::Value(ValueWithSpan { - value: Value::Null, - span: Span::empty(), - }), - ); - let Expr::Value(ValueWithSpan { - value: value @ Value::Placeholder(_), - .. - }) = to_wrap - else { - unreachable!("the Expr is known to be Expr::Value(ValueWithSpan::{{ value: Value::Placeholder(_), .. }})") - }; - - *expr = cast_to_v3_domain(value, &schema, &domain); - return Ok(true); - } - } - - Ok(false) - } - - fn would_edit(&mut self, node_path: &NodePath<'ast>, _target_node: &N) -> bool { - if let Some(( - node @ Expr::Value(ValueWithSpan { - value: Value::Placeholder(_), - .. - }), - )) = node_path.last_1_as() - { - if let Some(Type::Value(crate::unifier::Value::Eql(_))) = - self.node_types.get(&NodeKey::new(node)) - { - return true; - } - } - false - } -} diff --git a/packages/eql-mapper/src/transformation_rules/helpers.rs b/packages/eql-mapper/src/transformation_rules/helpers.rs index b8b8cf711..d6b8f8dda 100644 --- a/packages/eql-mapper/src/transformation_rules/helpers.rs +++ b/packages/eql-mapper/src/transformation_rules/helpers.rs @@ -1,28 +1,24 @@ +use std::collections::HashMap; +use std::mem; + use sqltk::parser::{ ast::{ BinaryOperator, CastKind, DataType, Expr, Function, FunctionArg, FunctionArgExpr, FunctionArgumentList, FunctionArguments, Ident, ObjectName, ObjectNamePart, + Value as SqltkValue, ValueWithSpan, }, tokenizer::Span, }; -use sqltk::NodePath; +use sqltk::NodeKey; -use crate::unifier::{DomainIdentity, EqlTerm}; +use crate::unifier::{EqlTerm, Type, Value}; -/// The v3 cast target `(schema, domain)` for a JSON ordering operand -/// ([`EqlTerm::JsonOrd`]) — always the shape-only scalar ord twin -/// `eql_v3.query_integer_ord`, regardless of the JSON leaf's scalar type. -/// -/// `eql_v3.ord_term` is type-agnostic (it extracts the `op` bytes as -/// `ope_cllw` and compares them bytewise), and `query_integer_ord`'s domain -/// CHECK is shape-only (`{v,i,op}`, no `c`). So a single twin serves numbers, -/// text, dates, etc. — which is what makes JSON range work in the extended -/// protocol, where the operand's scalar type is unknown at rewrite time. -/// Returns `None` for any other term. -/// The v3 cast target `(schema, domain)` for an encrypted-JSON *query operand* -/// whose domain is fixed by the operand's role rather than by the column's -/// domain identity. Returns `None` for any other term. +/// The v3 domain an encrypted **query operand** — the value side of a +/// predicate — casts to, or `None` if it takes no cast. /// +/// - [`EqlTerm::JsonAccessor`] / [`EqlTerm::JsonPath`] — no cast. A JSON field +/// selector is passed to the eql_v3 function as bare encrypted *text* +/// (`eql_v3."->"(json, text)`), not as a jsonb query payload. /// - [`EqlTerm::JsonOrd`] — the shape-only scalar ord twin /// `eql_v3.query_integer_ord`, regardless of the JSON leaf's scalar type. /// `eql_v3.ord_term` is type-agnostic (it extracts the `op` bytes as @@ -34,14 +30,37 @@ use crate::unifier::{DomainIdentity, EqlTerm}; /// needle domain. The fused value selector is a one-entry, term-less /// containment payload (`{sv: [{s}]}`), which is what /// `eql_v3.jsonb_contains` matches against. -pub(crate) fn json_query_operand_cast_target(eql_term: &EqlTerm) -> Option<(String, String)> { - let domain = match eql_term { - EqlTerm::JsonOrd(_) => "query_integer_ord", - EqlTerm::JsonValueSelector(_) => "query_json", - _ => return None, - }; +/// - Everything else — the column domain's `eql_v3.query_*` twin, which carries +/// only the search terms the predicate needs, not a whole ciphertext. +pub(crate) fn query_operand_domain(eql_term: &EqlTerm) -> Option<(String, String)> { + match eql_term { + EqlTerm::JsonAccessor(_) | EqlTerm::JsonPath(_) => None, + EqlTerm::JsonOrd(_) => Some(("eql_v3".to_string(), "query_integer_ord".to_string())), + EqlTerm::JsonValueSelector(_) => Some(("eql_v3".to_string(), "query_json".to_string())), + _ => { + let (schema, twin) = eql_term.eql_value().domain_identity().query_twin(); + Some((schema.to_string(), twin)) + } + } +} - Some(("eql_v3".to_string(), domain.to_string())) +/// The v3 domain an encrypted **full-payload** operand casts to: the column's +/// own domain, carrying the ciphertext plus every search term the column +/// indexes. +/// +/// This is what an `INSERT` value, an `UPDATE` assignment and a containment +/// needle all need — as opposed to a predicate operand, which needs only the +/// terms of [`query_operand_domain`]. +/// +/// Returns `None` for a JSON selector, which is bare text in every position. +pub(crate) fn full_payload_domain(eql_term: &EqlTerm) -> Option<(String, String)> { + match eql_term { + EqlTerm::JsonAccessor(_) | EqlTerm::JsonPath(_) => None, + _ => Some(( + "public".to_string(), + eql_term.eql_value().domain_identity().domain.value.clone(), + )), + } } /// The scalar comparison operators the v3 term-function rewrite handles. @@ -57,58 +76,51 @@ pub(crate) fn is_comparison_op(op: &BinaryOperator) -> bool { ) } -/// Whether an encrypted value at `node_path` is a **query operand** (the RHS of a -/// comparison or match predicate) rather than a **stored value** (an INSERT -/// `VALUES` item or UPDATE `SET` target). Walks the enclosing `Expr` ancestor -/// chain looking for a comparison `BinaryOp` or a `LIKE`/`ILIKE` predicate. The -/// traversal is post-order, so when a cast rule runs on the operand the enclosing -/// predicate is still intact in the path. -fn is_query_operand(node_path: &NodePath<'_>) -> bool { - let mut depth = 1; - while let Some(expr) = node_path.nth_last_as::(depth) { - match expr { - Expr::BinaryOp { op, .. } - if is_comparison_op(op) || matches!(op, BinaryOperator::AtAt) => - { - return true - } - Expr::Like { .. } | Expr::ILike { .. } => return true, - _ => {} - } - depth += 1; +/// Casts `target` — the already-transformed form of `original` — to the v3 +/// domain `domain_of` chooses for its EQL type. +/// +/// Only a literal or placeholder is cast. A column reference is already of its +/// domain type, and any other expression belongs to whichever rule owns it. +/// Returns `true` if a cast was applied. +/// +/// This is called by the rule that *owns the context* — a comparison, a match, +/// a containment, an INSERT value — so the choice of domain never has to be +/// inferred from where the node happens to sit in the tree. +pub(crate) fn cast_encrypted_operand( + node_types: &HashMap, Type>, + original: &Expr, + target: &mut Expr, + domain_of: fn(&EqlTerm) -> Option<(String, String)>, +) -> bool { + if !matches!(original, Expr::Value(_)) { + return false; } - false -} -/// The v3 cast target `(schema, domain typname)` for an encrypted value carrying -/// `identity` at `node_path`. A query operand casts to the `eql_v3.query_*` twin -/// (term-only payload); a stored value casts to the `public` column domain. -pub(crate) fn v3_cast_target( - node_path: &NodePath<'_>, - identity: &DomainIdentity, -) -> (String, String) { - if is_query_operand(node_path) { - let (schema, twin) = identity.query_twin(); - (schema.to_string(), twin) - } else { - ("public".to_string(), identity.domain.value.clone()) - } + let Some(Type::Value(Value::Eql(eql_term))) = node_types.get(&NodeKey::new(original)) else { + return false; + }; + + let Some((schema, domain)) = domain_of(eql_term) else { + return false; + }; + + let wrapped = mem::replace( + target, + Expr::Value(ValueWithSpan { + value: SqltkValue::Null, + span: Span::empty(), + }), + ); + + *target = cast_expr_to_v3_domain(wrapped, &schema, &domain); + true } -/// Builds `::JSONB::.` — the cast that wraps an encrypted -/// value (a jsonb payload) as an EQL v3 domain. `schema` is `public` for a stored -/// column domain and `eql_v3` for a query-operand twin. -pub(crate) fn cast_to_v3_domain( - wrapped: sqltk::parser::ast::Value, - schema: &str, - domain: &str, -) -> Expr { +/// Builds `::JSONB::.` around an arbitrary expression. +pub(crate) fn cast_expr_to_v3_domain(wrapped: Expr, schema: &str, domain: &str) -> Expr { let cast_jsonb = Expr::Cast { kind: CastKind::DoubleColon, - expr: Box::new(Expr::Value(sqltk::parser::ast::ValueWithSpan { - value: wrapped, - span: Span::empty(), - })), + expr: Box::new(wrapped), data_type: DataType::JSONB, format: None, }; diff --git a/packages/eql-mapper/src/transformation_rules/mod.rs b/packages/eql-mapper/src/transformation_rules/mod.rs index 4ac1f4007..82940fe99 100644 --- a/packages/eql-mapper/src/transformation_rules/mod.rs +++ b/packages/eql-mapper/src/transformation_rules/mod.rs @@ -11,8 +11,7 @@ mod helpers; -mod cast_literals_as_encrypted; -mod cast_params_as_encrypted; +mod cast_full_payload_operands; mod fail_on_placeholder_change; mod preserve_effective_aliases; mod rewrite_containment_ops; @@ -20,11 +19,11 @@ mod rewrite_eql_comparison_ops; mod rewrite_eql_match_ops; mod rewrite_json_value_selector_eq; mod rewrite_standard_sql_fns_on_eql_types; +mod substitute_encrypted_literals; use std::marker::PhantomData; -pub(crate) use cast_literals_as_encrypted::*; -pub(crate) use cast_params_as_encrypted::*; +pub(crate) use cast_full_payload_operands::*; pub(crate) use fail_on_placeholder_change::*; pub(crate) use preserve_effective_aliases::*; pub(crate) use rewrite_containment_ops::*; @@ -32,6 +31,7 @@ pub(crate) use rewrite_eql_comparison_ops::*; pub(crate) use rewrite_eql_match_ops::*; pub(crate) use rewrite_json_value_selector_eq::*; pub(crate) use rewrite_standard_sql_fns_on_eql_types::*; +pub(crate) use substitute_encrypted_literals::*; use crate::EqlMapperError; use sqltk::{NodePath, Transform, Visitable}; diff --git a/packages/eql-mapper/src/transformation_rules/rewrite_containment_ops.rs b/packages/eql-mapper/src/transformation_rules/rewrite_containment_ops.rs index 11b03deef..aa83a8b2e 100644 --- a/packages/eql-mapper/src/transformation_rules/rewrite_containment_ops.rs +++ b/packages/eql-mapper/src/transformation_rules/rewrite_containment_ops.rs @@ -13,6 +13,7 @@ use sqltk::{NodeKey, NodePath, Visitable}; use crate::unifier::{Type, Value}; use crate::EqlMapperError; +use super::helpers::{cast_encrypted_operand, full_payload_domain}; use super::TransformationRule; /// Rewrites JSON binary operators on encrypted columns to `eql_v3` function @@ -93,6 +94,17 @@ impl<'ast> TransformationRule<'ast> for RewriteContainmentOps<'ast> { target_node: &mut N, ) -> Result { if self.would_edit(node_path, target_node) { + // Read the original operands: `node_types` is keyed by them, and the + // cast this rule applies depends on which operator it is. + let Some((Expr::BinaryOp { + left: original_left, + right: original_right, + .. + },)) = node_path.last_1_as::() + else { + return Ok(false); + }; + let expr = target_node.downcast_mut::().unwrap(); if let Expr::BinaryOp { left, op, right } = expr { let fn_name = match op { @@ -103,8 +115,21 @@ impl<'ast> TransformationRule<'ast> for RewriteContainmentOps<'ast> { _ => return Ok(false), }; + // A containment needle is a whole encrypted document, so it + // casts to the column domain, not to a query twin. A `->`/`->>` + // selector takes no cast at all — `full_payload_domain` returns + // `None` for it — because `eql_v3."->"(json, text)` wants the + // bare encrypted selector text. + cast_encrypted_operand(&self.node_types, original_left, left, full_payload_domain); + cast_encrypted_operand( + &self.node_types, + original_right, + right, + full_payload_domain, + ); + // Use mem::replace to move (not copy) the original nodes, - // preserving their NodeKey identity for downstream casting rules + // preserving their NodeKey identity for downstream rules let dummy = Expr::Value(ValueWithSpan { value: SqltkValue::Null, span: Span::empty(), diff --git a/packages/eql-mapper/src/transformation_rules/rewrite_eql_comparison_ops.rs b/packages/eql-mapper/src/transformation_rules/rewrite_eql_comparison_ops.rs index 258e527f1..acb2f8d47 100644 --- a/packages/eql-mapper/src/transformation_rules/rewrite_eql_comparison_ops.rs +++ b/packages/eql-mapper/src/transformation_rules/rewrite_eql_comparison_ops.rs @@ -10,7 +10,9 @@ use sqltk::{NodeKey, NodePath, Visitable}; use crate::unifier::{DomainIdentity, EqlTerm, Type, Value}; use crate::EqlMapperError; -use super::helpers::{eql_v3_term_call, is_comparison_op}; +use super::helpers::{ + cast_encrypted_operand, eql_v3_term_call, is_comparison_op, query_operand_domain, +}; use super::TransformationRule; /// Rewrites scalar comparison operators on encrypted columns into the EQL v3 @@ -109,15 +111,26 @@ impl<'ast> TransformationRule<'ast> for RewriteEqlComparisonOps<'ast> { ))); }; - if let Expr::BinaryOp { left, right, .. } = target_node.downcast_mut::().unwrap() { + if let Expr::BinaryOp { + left: target_left, + right: target_right, + .. + } = target_node.downcast_mut::().unwrap() + { + // Cast the operands before wrapping them: this rule owns the + // comparison, so it knows both are query operands and casts them to + // the term-only `eql_v3.query_*` twin. + cast_encrypted_operand(&self.node_types, left, target_left, query_operand_domain); + cast_encrypted_operand(&self.node_types, right, target_right, query_operand_domain); + let dummy = Expr::Value(ValueWithSpan { value: SqltkValue::Null, span: Span::empty(), }); - let left_expr = mem::replace(&mut **left, dummy.clone()); - let right_expr = mem::replace(&mut **right, dummy); - **left = eql_v3_term_call(term_fn, left_expr); - **right = eql_v3_term_call(term_fn, right_expr); + let left_expr = mem::replace(&mut **target_left, dummy.clone()); + let right_expr = mem::replace(&mut **target_right, dummy); + **target_left = eql_v3_term_call(term_fn, left_expr); + **target_right = eql_v3_term_call(term_fn, right_expr); return Ok(true); } diff --git a/packages/eql-mapper/src/transformation_rules/rewrite_eql_match_ops.rs b/packages/eql-mapper/src/transformation_rules/rewrite_eql_match_ops.rs index 671acbde6..ca9a1a4f0 100644 --- a/packages/eql-mapper/src/transformation_rules/rewrite_eql_match_ops.rs +++ b/packages/eql-mapper/src/transformation_rules/rewrite_eql_match_ops.rs @@ -10,7 +10,7 @@ use sqltk::{NodeKey, NodePath, Visitable}; use crate::unifier::{DomainIdentity, Type, Value}; use crate::EqlMapperError; -use super::helpers::eql_v3_term_call; +use super::helpers::{cast_encrypted_operand, eql_v3_term_call, query_operand_domain}; use super::TransformationRule; /// Rewrites fuzzy-match predicates on encrypted columns into the EQL v3 @@ -92,20 +92,44 @@ impl<'ast> TransformationRule<'ast> for RewriteEqlMatchOps<'ast> { ))); }; + // The pattern is a query operand of this predicate — cast it to the + // column's `eql_v3.query_*` twin before wrapping. This rule owns the + // predicate, so no ancestor inspection is needed to know that. + let (original_col, original_pat) = match original { + Expr::Like { expr, pattern, .. } | Expr::ILike { expr, pattern, .. } => { + (&**expr, &**pattern) + } + Expr::BinaryOp { left, right, .. } => (&**left, &**right), + _ => return Ok(false), + }; + let dummy = Expr::Value(ValueWithSpan { value: SqltkValue::Null, span: Span::empty(), }); let expr = target_node.downcast_mut::().unwrap(); let (col_expr, pat_expr) = match expr { - Expr::Like { expr, pattern, .. } | Expr::ILike { expr, pattern, .. } => ( - mem::replace(&mut **expr, dummy.clone()), - mem::replace(&mut **pattern, dummy), - ), - Expr::BinaryOp { left, right, .. } => ( - mem::replace(&mut **left, dummy.clone()), - mem::replace(&mut **right, dummy), - ), + Expr::Like { expr, pattern, .. } | Expr::ILike { expr, pattern, .. } => { + cast_encrypted_operand(&self.node_types, original_col, expr, query_operand_domain); + cast_encrypted_operand( + &self.node_types, + original_pat, + pattern, + query_operand_domain, + ); + ( + mem::replace(&mut **expr, dummy.clone()), + mem::replace(&mut **pattern, dummy), + ) + } + Expr::BinaryOp { left, right, .. } => { + cast_encrypted_operand(&self.node_types, original_col, left, query_operand_domain); + cast_encrypted_operand(&self.node_types, original_pat, right, query_operand_domain); + ( + mem::replace(&mut **left, dummy.clone()), + mem::replace(&mut **right, dummy), + ) + } _ => return Ok(false), }; diff --git a/packages/eql-mapper/src/transformation_rules/rewrite_json_value_selector_eq.rs b/packages/eql-mapper/src/transformation_rules/rewrite_json_value_selector_eq.rs index 8af09435f..52a873e0e 100644 --- a/packages/eql-mapper/src/transformation_rules/rewrite_json_value_selector_eq.rs +++ b/packages/eql-mapper/src/transformation_rules/rewrite_json_value_selector_eq.rs @@ -13,6 +13,7 @@ use sqltk::{NodeKey, NodePath, Visitable}; use crate::unifier::{EqlTerm, Type, Value}; use crate::EqlMapperError; +use super::helpers::{cast_encrypted_operand, query_operand_domain}; use super::TransformationRule; /// Rewrites equality on an encrypted JSON **field** into value-selector @@ -160,15 +161,25 @@ impl<'ast> TransformationRule<'ast> for RewriteJsonValueSelectorEq<'ast> { // Move (not clone) the transformed value operand so its NodeKey identity // survives; the container comes from the original AST, where it is still // the bare column reference the containment call needs. + // The needle is a query operand of this rule's own predicate: cast it to + // `eql_v3.query_json`, the containment-needle domain. + let (original_needle, target_needle) = if value_on_right { + (right, target_right) + } else { + (left, target_left) + }; + cast_encrypted_operand( + &self.node_types, + original_needle, + target_needle, + query_operand_domain, + ); + let dummy = Expr::Value(ValueWithSpan { value: SqltkValue::Null, span: Span::empty(), }); - let needle = if value_on_right { - mem::replace(&mut **target_right, dummy) - } else { - mem::replace(&mut **target_left, dummy) - }; + let needle = mem::replace(&mut **target_needle, dummy); let contains = Self::jsonb_contains(container.clone(), needle); diff --git a/packages/eql-mapper/src/transformation_rules/substitute_encrypted_literals.rs b/packages/eql-mapper/src/transformation_rules/substitute_encrypted_literals.rs new file mode 100644 index 000000000..0e1281f34 --- /dev/null +++ b/packages/eql-mapper/src/transformation_rules/substitute_encrypted_literals.rs @@ -0,0 +1,74 @@ +use std::{any::type_name, collections::HashMap}; + +use sqltk::parser::ast::{Expr, Value, ValueWithSpan}; +use sqltk::parser::tokenizer::Span; +use sqltk::{NodeKey, NodePath, Visitable}; + +use crate::EqlMapperError; + +use super::TransformationRule; + +/// Replaces each plaintext literal with the encrypted value the proxy produced +/// for it. +/// +/// Substitution only — **no casting**. Where a literal appears determines the +/// domain it must be cast to, and that is known with certainty only by the rule +/// that owns the surrounding construct (a comparison, a containment, an +/// `INSERT`). Those rules apply the cast; this one just puts the ciphertext in +/// place, wherever it is. +/// +/// Runs before the casting rules in the tuple so the value is already in place +/// by the time a rule wraps it. +#[derive(Debug)] +pub struct SubstituteEncryptedLiterals<'ast> { + encrypted_literals: HashMap, Value>, +} + +impl<'ast> SubstituteEncryptedLiterals<'ast> { + pub fn new(encrypted_literals: HashMap, Value>) -> Self { + Self { encrypted_literals } + } +} + +impl<'ast> TransformationRule<'ast> for SubstituteEncryptedLiterals<'ast> { + fn apply( + &mut self, + node_path: &NodePath<'ast>, + target_node: &mut N, + ) -> Result { + let Some((Expr::Value(ValueWithSpan { value, .. }),)) = node_path.last_1_as::() + else { + return Ok(false); + }; + + let Some(replacement) = self.encrypted_literals.remove(&NodeKey::new(value)) else { + return Ok(false); + }; + + let target_node = target_node.downcast_mut::().unwrap(); + *target_node = Expr::Value(ValueWithSpan { + value: replacement, + span: Span::empty(), + }); + + Ok(true) + } + + fn would_edit(&mut self, node_path: &NodePath<'ast>, _target_node: &N) -> bool { + if let Some((Expr::Value(ValueWithSpan { value, .. }),)) = node_path.last_1_as::() { + return self.encrypted_literals.contains_key(&NodeKey::new(value)); + } + false + } + + fn check_postcondition(&self) -> Result<(), EqlMapperError> { + if self.encrypted_literals.is_empty() { + Ok(()) + } else { + Err(EqlMapperError::Transform(format!( + "Postcondition failed in {}: unused encrypted literals", + type_name::() + ))) + } + } +} diff --git a/packages/eql-mapper/src/type_checked_statement.rs b/packages/eql-mapper/src/type_checked_statement.rs index ef12d807e..c86590efb 100644 --- a/packages/eql-mapper/src/type_checked_statement.rs +++ b/packages/eql-mapper/src/type_checked_statement.rs @@ -5,10 +5,10 @@ use sqltk::{AsNodeKey, NodeKey, Transformable}; use crate::unifier::EqlTerm; use crate::{ - CastLiteralsAsEncrypted, CastParamsAsEncrypted, DryRunnable, EqlMapperError, - FailOnPlaceholderChange, JsonValueSelectors, OutputParam, OutputParamSource, Param, ParamPlan, - PreserveEffectiveAliases, RenumberParams, RewriteContainmentOps, RewriteEqlComparisonOps, - RewriteEqlMatchOps, RewriteJsonValueSelectorEq, RewriteStandardSqlFnsOnEqlTypes, + CastFullPayloadOperands, DryRunnable, EqlMapperError, FailOnPlaceholderChange, + JsonValueSelectors, OutputParam, OutputParamSource, Param, ParamPlan, PreserveEffectiveAliases, + RenumberParams, RewriteContainmentOps, RewriteEqlComparisonOps, RewriteEqlMatchOps, + RewriteJsonValueSelectorEq, RewriteStandardSqlFnsOnEqlTypes, SubstituteEncryptedLiterals, TransformationRule, }; @@ -249,16 +249,20 @@ impl<'ast> TypeCheckedStatement<'ast> { &self, encrypted_literals: HashMap, sqltk::parser::ast::Value>, ) -> DryRunnable<'_, impl TransformationRule<'_>> { + // Substitution runs first so every encrypted literal is in place before + // any rule wraps it. Each rewrite rule then applies the cast its own + // context requires, and `CastFullPayloadOperands` covers the one + // context that has no rewrite of its own — INSERT and UPDATE values. DryRunnable::new(( + SubstituteEncryptedLiterals::new(encrypted_literals), RewriteStandardSqlFnsOnEqlTypes::new(Arc::clone(&self.node_types)), RewriteContainmentOps::new(Arc::clone(&self.node_types)), RewriteJsonValueSelectorEq::new(Arc::clone(&self.node_types)), RewriteEqlComparisonOps::new(Arc::clone(&self.node_types)), RewriteEqlMatchOps::new(Arc::clone(&self.node_types)), + CastFullPayloadOperands::new(Arc::clone(&self.node_types)), PreserveEffectiveAliases, - CastLiteralsAsEncrypted::new(encrypted_literals, Arc::clone(&self.node_types)), FailOnPlaceholderChange::new(), - CastParamsAsEncrypted::new(Arc::clone(&self.node_types)), )) } } From c1cc1721a603916653bb3dc6bf90ed254f7562ac Mon Sep 17 00:00:00 2001 From: James Sadler Date: Sun, 26 Jul 2026 11:23:10 +1000 Subject: [PATCH 06/10] fix(mapper): ORDER BY an encrypted column must order by its ordering term MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ORDER BY col` on an encrypted column was passed through untouched. The column is a domain over `jsonb`, so PostgreSQL compared whole payloads by jsonb rules — and jsonb compares objects field by field starting at `c`, the ciphertext, which is randomised per encryption. Results came back in an order that was not merely wrong but different on every insert. Rewrites to the ordering term, which is what makes the comparison meaningful: ORDER BY col -> ORDER BY eql_v3.ord_term(col) ORDER BY t.col DESC -> ORDER BY eql_v3.ord_term(t.col) DESC ORDER BY col NULLS FIRST-> ORDER BY eql_v3.ord_term(col) NULLS FIRST (`ord_term_ore` for block-ORE domains.) Sort options are untouched, so ASC/DESC and NULLS FIRST/LAST keep working — the term is a plain orderable value and NULL stays NULL. `ord_term` yields `ope_cllw`, a `bytea` domain ordered bytewise; `ord_term_ore` yields `ore_block_256`, ordered by its own btree operator class. Verified directly against PostgreSQL: over four values re-encrypted five times, the bare `ORDER BY` gave a different wrong order every run while the term ordering was correct every run. Ordering by a column whose domain carries no ordering term is now a capability error rather than an arbitrary sort, matching how `RewriteEqlComparisonOps` treats an unsupported operator. This is the cause of the "flaky" ORE/OPE ordering tests. They were not flaky: the ordering was random, so a test comparing two values passed about half the time while one comparing five values essentially never did. Integration: 236 -> 298 passing, 168 -> 106 failing, nothing newly failing. Fixes all of map_ope_index_order (6), map_ore_index_order (18) and select::order_by (42), and those families are now deterministic across repeated runs. Mapper units 104 -> 107. Stable-Commit-Id: q-7znx6wzcc63yyk8hnganvrbdce --- packages/eql-mapper/src/lib.rs | 86 +++++++++++++++ .../src/transformation_rules/mod.rs | 2 + .../rewrite_eql_order_by.rs | 101 ++++++++++++++++++ .../eql-mapper/src/type_checked_statement.rs | 5 +- 4 files changed, 192 insertions(+), 2 deletions(-) create mode 100644 packages/eql-mapper/src/transformation_rules/rewrite_eql_order_by.rs diff --git a/packages/eql-mapper/src/lib.rs b/packages/eql-mapper/src/lib.rs index eae3c93eb..a4f6e071f 100644 --- a/packages/eql-mapper/src/lib.rs +++ b/packages/eql-mapper/src/lib.rs @@ -2537,6 +2537,92 @@ mod test { ); } + /// `ORDER BY` on an encrypted column must order by its ordering TERM. A bare + /// `ORDER BY col` compares the jsonb payloads, whose first field is the + /// randomised ciphertext — silently returning rows in an arbitrary order + /// that differs on every insert. + #[test] + fn order_by_encrypted_column_uses_ord_term() { + let schema = resolver(schema! { + tables: { + employees: { + id, + salary (EQL("eql_v3_integer_ord"): Ord), + } + } + }); + + for (order_by, expected) in [ + ("salary", "eql_v3.ord_term(salary)"), + ("salary DESC", "eql_v3.ord_term(salary) DESC"), + ( + "salary ASC NULLS FIRST", + "eql_v3.ord_term(salary) ASC NULLS FIRST", + ), + ( + "employees.salary DESC NULLS LAST", + "eql_v3.ord_term(employees.salary) DESC NULLS LAST", + ), + // A native column is left alone; a mixed list rewrites only the + // encrypted term. + ("id", "id"), + ("id, salary", "id, eql_v3.ord_term(salary)"), + ] { + let statement = parse(&format!("SELECT id FROM employees ORDER BY {order_by}")); + let typed = type_check(schema.clone(), &statement).unwrap(); + + assert_eq!( + typed.transform(HashMap::new()).unwrap().to_string(), + format!("SELECT id FROM employees ORDER BY {expected}"), + "unexpected rewrite for `ORDER BY {order_by}`" + ); + } + } + + /// A block-ORE domain orders through `ord_term_ore`. + #[test] + fn order_by_ore_column_uses_ord_term_ore() { + let schema = resolver(schema! { + tables: { + employees: { + id, + salary (EQL("eql_v3_integer_ord_ore"): Ord), + } + } + }); + + let statement = parse("SELECT id FROM employees ORDER BY salary"); + let typed = type_check(schema, &statement).unwrap(); + + assert_eq!( + typed.transform(HashMap::new()).unwrap().to_string(), + "SELECT id FROM employees ORDER BY eql_v3.ord_term_ore(salary)" + ); + } + + /// Ordering by a column whose domain carries no ordering term is a + /// capability error, not an arbitrary sort. + #[test] + fn order_by_column_without_ordering_term_is_an_error() { + let schema = resolver(schema! { + tables: { + employees: { + id, + name (EQL("eql_v3_text_match"): TokenMatch), + } + } + }); + + let statement = parse("SELECT id FROM employees ORDER BY name"); + let typed = type_check(schema, &statement).unwrap(); + + let err = typed.transform(HashMap::new()).unwrap_err().to_string(); + assert!( + err.contains("ORDER BY") && err.contains("no ordering term"), + "expected a capability error, got: {err}" + ); + } + #[test] fn jsonb_path_query_param_to_eql() { // init_tracing(); diff --git a/packages/eql-mapper/src/transformation_rules/mod.rs b/packages/eql-mapper/src/transformation_rules/mod.rs index 82940fe99..f3ff77fcc 100644 --- a/packages/eql-mapper/src/transformation_rules/mod.rs +++ b/packages/eql-mapper/src/transformation_rules/mod.rs @@ -17,6 +17,7 @@ mod preserve_effective_aliases; mod rewrite_containment_ops; mod rewrite_eql_comparison_ops; mod rewrite_eql_match_ops; +mod rewrite_eql_order_by; mod rewrite_json_value_selector_eq; mod rewrite_standard_sql_fns_on_eql_types; mod substitute_encrypted_literals; @@ -29,6 +30,7 @@ pub(crate) use preserve_effective_aliases::*; pub(crate) use rewrite_containment_ops::*; pub(crate) use rewrite_eql_comparison_ops::*; pub(crate) use rewrite_eql_match_ops::*; +pub(crate) use rewrite_eql_order_by::*; pub(crate) use rewrite_json_value_selector_eq::*; pub(crate) use rewrite_standard_sql_fns_on_eql_types::*; pub(crate) use substitute_encrypted_literals::*; diff --git a/packages/eql-mapper/src/transformation_rules/rewrite_eql_order_by.rs b/packages/eql-mapper/src/transformation_rules/rewrite_eql_order_by.rs new file mode 100644 index 000000000..b1a1907ad --- /dev/null +++ b/packages/eql-mapper/src/transformation_rules/rewrite_eql_order_by.rs @@ -0,0 +1,101 @@ +use std::collections::HashMap; +use std::mem; +use std::sync::Arc; + +use sqltk::parser::ast::Value as SqltkValue; +use sqltk::parser::ast::{Expr, OrderByExpr, ValueWithSpan}; +use sqltk::parser::tokenizer::Span; +use sqltk::{NodeKey, NodePath, Visitable}; + +use crate::unifier::{DomainIdentity, Type, Value}; +use crate::EqlMapperError; + +use super::helpers::eql_v3_term_call; +use super::TransformationRule; + +/// Rewrites `ORDER BY` on an encrypted column to order by its **ordering term**: +/// +/// - `ORDER BY col` → `ORDER BY eql_v3.ord_term(col)` +/// - `ORDER BY t.col DESC` → `ORDER BY eql_v3.ord_term(t.col) DESC` +/// - `ORDER BY col NULLS FIRST`→ `ORDER BY eql_v3.ord_term(col) NULLS FIRST` +/// +/// (`ord_term_ore` for block-ORE domains.) Sort options are untouched, so +/// `ASC`/`DESC` and `NULLS FIRST`/`LAST` keep working — the term is a plain +/// orderable value and `NULL` stays `NULL`. +/// +/// **Without this rewrite the results are silently misordered.** An encrypted +/// column is a domain over `jsonb`, so a bare `ORDER BY` compares whole payloads +/// by jsonb rules — and jsonb compares objects field by field starting at `c`, +/// the ciphertext, which is randomised per encryption. The rows come back in an +/// order that is not just wrong but *different on every insert*. The ordering +/// terms exist precisely to avoid that: `ord_term` yields `ope_cllw` (a `bytea` +/// domain, ordered bytewise) and `ord_term_ore` yields `ore_block_256` (ordered +/// by its own btree operator class). +/// +/// A column whose domain carries no ordering term cannot be ordered at all, and +/// is a capability error rather than a silent arbitrary sort. +#[derive(Debug)] +pub struct RewriteEqlOrderBy<'ast> { + node_types: Arc, Type>>, +} + +impl<'ast> RewriteEqlOrderBy<'ast> { + pub fn new(node_types: Arc, Type>>) -> Self { + Self { node_types } + } + + fn eql_identity_of(&self, expr: &'ast Expr) -> Option { + match self.node_types.get(&NodeKey::new(expr)) { + Some(Type::Value(Value::Eql(eql_term))) => { + Some(eql_term.eql_value().domain_identity().clone()) + } + _ => None, + } + } +} + +impl<'ast> TransformationRule<'ast> for RewriteEqlOrderBy<'ast> { + fn apply( + &mut self, + node_path: &NodePath<'ast>, + target_node: &mut N, + ) -> Result { + // Read the identity from the ORIGINAL node — `node_types` is keyed by it. + let Some((original,)) = node_path.last_1_as::() else { + return Ok(false); + }; + + let Some(identity) = self.eql_identity_of(&original.expr) else { + return Ok(false); + }; + + let Some(term_fn) = identity.ord_term_fn() else { + return Err(EqlMapperError::Transform(format!( + "encrypted column {} cannot be used in ORDER BY (domain {} carries no ordering term)", + identity.token, identity.domain.value + ))); + }; + + let Some(target) = target_node.downcast_mut::() else { + return Ok(false); + }; + + let expr = mem::replace( + &mut target.expr, + Expr::Value(ValueWithSpan { + value: SqltkValue::Null, + span: Span::empty(), + }), + ); + target.expr = eql_v3_term_call(term_fn, expr); + + Ok(true) + } + + fn would_edit(&mut self, node_path: &NodePath<'ast>, _target_node: &N) -> bool { + match node_path.last_1_as::() { + Some((original,)) => self.eql_identity_of(&original.expr).is_some(), + None => false, + } + } +} diff --git a/packages/eql-mapper/src/type_checked_statement.rs b/packages/eql-mapper/src/type_checked_statement.rs index c86590efb..0209a5f64 100644 --- a/packages/eql-mapper/src/type_checked_statement.rs +++ b/packages/eql-mapper/src/type_checked_statement.rs @@ -8,8 +8,8 @@ use crate::{ CastFullPayloadOperands, DryRunnable, EqlMapperError, FailOnPlaceholderChange, JsonValueSelectors, OutputParam, OutputParamSource, Param, ParamPlan, PreserveEffectiveAliases, RenumberParams, RewriteContainmentOps, RewriteEqlComparisonOps, RewriteEqlMatchOps, - RewriteJsonValueSelectorEq, RewriteStandardSqlFnsOnEqlTypes, SubstituteEncryptedLiterals, - TransformationRule, + RewriteEqlOrderBy, RewriteJsonValueSelectorEq, RewriteStandardSqlFnsOnEqlTypes, + SubstituteEncryptedLiterals, TransformationRule, }; use crate::unifier::{Projection, Type, Value}; @@ -260,6 +260,7 @@ impl<'ast> TypeCheckedStatement<'ast> { RewriteJsonValueSelectorEq::new(Arc::clone(&self.node_types)), RewriteEqlComparisonOps::new(Arc::clone(&self.node_types)), RewriteEqlMatchOps::new(Arc::clone(&self.node_types)), + RewriteEqlOrderBy::new(Arc::clone(&self.node_types)), CastFullPayloadOperands::new(Arc::clone(&self.node_types)), PreserveEffectiveAliases, FailOnPlaceholderChange::new(), From 7259d6b2a2ab4e605cfb9007d08a92780520dd69 Mon Sep 17 00:00:00 2001 From: James Sadler Date: Sun, 26 Jul 2026 12:06:51 +1000 Subject: [PATCH 07/10] fix(proxy): project encrypted query operands with into_query_operand MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every scalar predicate on an encrypted column failed: ERROR: value for domain eql_v3.query_integer_ord violates check constraint "query_integer_ord_check" A query operand carries the column's search terms but never a decryptable ciphertext — the `eql_v3.query_*` domains enforce that with `NOT (VALUE ? 'c')`. The proxy encrypted every scalar operand with `EqlOperation::Store` and sent the storage payload, `c` and all, into a query position. It cannot simply encrypt differently: a single `EqlOperation::Query` yields only ONE term, while an operand may need several (`{v,i,hm,ob}` for `query_text_ord`). The client's answer is to encrypt in Store mode and project — `EqlCiphertextV3::into_query_operand`, whose own test is named `into_query_operand_scalar_keeps_all_terms_and_drops_c`. The proxy never called it. Projecting needs to know which operands are query operands, and nothing carried that: the operand of `WHERE col = $1` and the value of `INSERT ... VALUES ($1)` are both `EqlTerm::Full`. So the mapper now records it, in the new `QueryOperands`, alongside the existing `JsonValueSelectors`. Membership is decided syntactically by the predicate an operand belongs to — comparisons, `LIKE`/`ILIKE`, `@@` — which is the same set of contexts whose rewrite rules cast to a `eql_v3.query_*` twin. Containment (`@>`/`<@`) is deliberately excluded: its needle is a whole document and keeps its full payload, as do INSERT values and UPDATE assignments. The flag rides through `ParamPlan` to the proxy's `OutputParam`, so Bind projects the params it sends, and the literal path projects at Parse. Payloads the encryptor already produced query-shaped (the JSON selectors and SteVec terms, which do go through `EqlOperation::Query`) are left alone. Fixes map_ope_index_where (6), map_ore_index_where (6), map_match_index (1) and 7 of 8 map_unique_index — scalar encrypted search worked in neither protocol before this. Integration 298 -> 317 passing, 106 -> 87 failing, nothing newly failing. The remaining map_unique_index failure (`..._all_with_wildcard`, a `SELECT *` with six encrypted params) fails with a protocol-level UnexpectedMessage and is a separate, pre-existing bug. Stable-Commit-Id: q-4jyyztp3wxv6ve8ndnwpp59r2y --- .../src/postgresql/context/statement.rs | 5 ++ .../src/postgresql/frontend.rs | 40 +++++++++++- .../src/postgresql/messages/parse.rs | 3 + packages/eql-mapper/src/eql_mapper.rs | 1 + .../src/inference/infer_type_impls/expr.rs | 40 ++++++++++++ packages/eql-mapper/src/inference/mod.rs | 23 ++++++- packages/eql-mapper/src/lib.rs | 2 + packages/eql-mapper/src/param_plan.rs | 5 ++ packages/eql-mapper/src/query_operands.rs | 61 +++++++++++++++++++ .../eql-mapper/src/type_checked_statement.rs | 12 ++++ 10 files changed, 189 insertions(+), 3 deletions(-) create mode 100644 packages/eql-mapper/src/query_operands.rs diff --git a/packages/cipherstash-proxy/src/postgresql/context/statement.rs b/packages/cipherstash-proxy/src/postgresql/context/statement.rs index 920897361..77053a59b 100644 --- a/packages/cipherstash-proxy/src/postgresql/context/statement.rs +++ b/packages/cipherstash-proxy/src/postgresql/context/statement.rs @@ -48,6 +48,10 @@ pub struct OutputParam { /// The input param(s) its value is built from. pub source: OutputParamSource, + + /// Whether this param is a query operand, whose payload must be projected + /// to carry search terms without a ciphertext. + pub query_operand: bool, } /// @@ -128,6 +132,7 @@ pub fn output_params_from_plan( .zip(output_columns) .map(|(output, column)| OutputParam { column, + query_operand: output.query_operand, source: match &output.source { eql_mapper::OutputParamSource::Input(param) => { OutputParamSource::Input(to_index(param.0)) diff --git a/packages/cipherstash-proxy/src/postgresql/frontend.rs b/packages/cipherstash-proxy/src/postgresql/frontend.rs index 3de073628..24cda6dbf 100644 --- a/packages/cipherstash-proxy/src/postgresql/frontend.rs +++ b/packages/cipherstash-proxy/src/postgresql/frontend.rs @@ -609,7 +609,7 @@ where let start = Instant::now(); - let encrypted = self + let mut encrypted = self .context .encrypt(plaintexts, literal_columns) .await @@ -617,6 +617,13 @@ where counter!(ENCRYPTION_ERROR_TOTAL).increment(1); })?; + for ((_, literal), encrypted) in literal_values.iter().zip(encrypted.iter_mut()) { + project_query_operand( + typed_statement.query_operands.contains_literal(literal), + encrypted, + ); + } + debug!(target: MAPPER, client_id = self.context.client_id, ?literal_columns, @@ -952,6 +959,7 @@ where .map(|(idx, column)| OutputParam { column: column.to_owned(), source: OutputParamSource::Input(idx), + query_operand: false, }) .collect(); @@ -1096,7 +1104,7 @@ where let start = Instant::now(); - let encrypted = self + let mut encrypted = self .context .encrypt(plaintexts, &output_param_columns) .await @@ -1104,6 +1112,10 @@ where counter!(ENCRYPTION_ERROR_TOTAL).increment(1); })?; + for (output, encrypted) in statement.output_params.iter().zip(encrypted.iter_mut()) { + project_query_operand(output.query_operand, encrypted); + } + let duration = Instant::now().duration_since(start); // Record timing and metadata for this encryption operation @@ -1204,6 +1216,30 @@ where } } +/// Projects a stored payload into its query operand when the value is bound in +/// a predicate rather than stored. +/// +/// A query operand carries the column's search terms but never a decryptable +/// ciphertext — the `eql_v3.query_*` domains reject `c` outright. It cannot be +/// produced by encrypting differently: a single `EqlOperation::Query` yields +/// only ONE term, while an operand may need several (`{v,i,hm,ob}`), so the +/// value is encrypted in Store mode and projected here. +/// +/// Terms that are already query-shaped (the JSON selectors and SteVec terms, +/// which the encryptor produces via `EqlOperation::Query`) are left alone. +fn project_query_operand(query_operand: bool, encrypted: &mut Option) { + if !query_operand { + return; + } + + match encrypted.take() { + Some(EqlOutput::Store(ciphertext)) => { + *encrypted = Some(EqlOutput::Query(ciphertext.into_query_operand())); + } + already_query_shaped => *encrypted = already_query_shaped, + } +} + fn literals_to_plaintext( typed_statement: &TypeCheckedStatement<'_>, literal_columns: &Vec>, diff --git a/packages/cipherstash-proxy/src/postgresql/messages/parse.rs b/packages/cipherstash-proxy/src/postgresql/messages/parse.rs index e356b54bf..e0f4ec62f 100644 --- a/packages/cipherstash-proxy/src/postgresql/messages/parse.rs +++ b/packages/cipherstash-proxy/src/postgresql/messages/parse.rs @@ -191,10 +191,12 @@ mod tests { OutputParam { column: None, source: OutputParamSource::Input(0), + query_operand: false, }, OutputParam { column: Some(column), source: OutputParamSource::Input(1), + query_operand: false, }, ]; @@ -225,6 +227,7 @@ mod tests { let output_params = vec![OutputParam { column: None, source: OutputParamSource::Input(1), + query_operand: false, }]; parse.rewrite_param_types(&output_params); diff --git a/packages/eql-mapper/src/eql_mapper.rs b/packages/eql-mapper/src/eql_mapper.rs index b2c23ceb7..2df75aec1 100644 --- a/packages/eql-mapper/src/eql_mapper.rs +++ b/packages/eql-mapper/src/eql_mapper.rs @@ -181,6 +181,7 @@ impl<'ast> EqlMapper<'ast> { params, literals, self.inferencer.borrow().take_json_value_selectors(), + self.inferencer.borrow().take_query_operands(), Arc::new(node_types), )) } diff --git a/packages/eql-mapper/src/inference/infer_type_impls/expr.rs b/packages/eql-mapper/src/inference/infer_type_impls/expr.rs index 9b92b9266..8a8524031 100644 --- a/packages/eql-mapper/src/inference/infer_type_impls/expr.rs +++ b/packages/eql-mapper/src/inference/infer_type_impls/expr.rs @@ -190,6 +190,24 @@ impl<'ast> InferType<'ast, Expr> for TypeInferencer<'ast> { if !handled { get_sql_binop_rule(op).apply_constraints(self, left, right, expr_val)?; } + + // The operands of a predicate reach PostgreSQL as query + // operands — terms only, never a ciphertext. Record them so the + // proxy projects their payloads accordingly. Containment + // (`@>`/`<@`) is deliberately excluded: its needle is a whole + // document and keeps its full payload. + if matches!( + op, + BinaryOperator::Eq + | BinaryOperator::NotEq + | BinaryOperator::Lt + | BinaryOperator::LtEq + | BinaryOperator::Gt + | BinaryOperator::GtEq + | BinaryOperator::AtAt + ) { + self.record_query_operands([&**left, &**right]); + } } // `customer_name LIKE 'A%'`. Route LIKE/ILIKE through the `~~`/`~~*` @@ -210,6 +228,7 @@ impl<'ast> InferType<'ast, Expr> for TypeInferencer<'ast> { BinaryOperator::PGLikeMatch }; get_sql_binop_rule(&op).apply_constraints(self, expr, pattern, expr_val)?; + self.record_query_operands([&**expr, &**pattern]); } Expr::ILike { negated, @@ -224,6 +243,7 @@ impl<'ast> InferType<'ast, Expr> for TypeInferencer<'ast> { BinaryOperator::PGILikeMatch }; get_sql_binop_rule(&op).apply_constraints(self, expr, pattern, expr_val)?; + self.record_query_operands([&**expr, &**pattern]); } Expr::Like { any: true, .. } | Expr::ILike { any: true, .. } => { @@ -590,6 +610,26 @@ impl<'ast> TypeInferencer<'ast> { self.eql_json_value(expr).map(|json| (json, selector)) } + /// Records each of `exprs` that is a literal or placeholder as a query + /// operand — an operand of a predicate, which reaches PostgreSQL carrying + /// only search terms and never a ciphertext. + /// + /// Column references are ignored: they are already stored payloads, and it + /// is only the bound values whose encryption shape this decides. + fn record_query_operands(&self, exprs: impl IntoIterator) { + for expr in exprs { + match Self::as_ast_value(expr) { + Some(ast::Value::Placeholder(placeholder)) => { + if let Ok(param) = Param::try_from(placeholder) { + self.record_query_operand_param(param); + } + } + Some(node) => self.record_query_operand_literal(node), + None => {} + } + } + } + /// Types `value` — the value half of `col -> sel = value` — as a fused /// value selector, and records where its path half (`selector`) comes from. /// diff --git a/packages/eql-mapper/src/inference/mod.rs b/packages/eql-mapper/src/inference/mod.rs index 58305159c..dc184593a 100644 --- a/packages/eql-mapper/src/inference/mod.rs +++ b/packages/eql-mapper/src/inference/mod.rs @@ -19,7 +19,8 @@ use sqltk::parser::ast::{ use sqltk::{into_control_flow, AsNodeKey, Break, Visitable, Visitor}; use crate::{ - JsonSelectorSource, JsonValueSelectors, Param, ScopeError, ScopeTracker, TableResolver, + JsonSelectorSource, JsonValueSelectors, Param, QueryOperands, ScopeError, ScopeTracker, + TableResolver, }; pub(crate) use registry::*; @@ -59,6 +60,12 @@ pub struct TypeInferencer<'ast> { /// the path for which value ([`crate::JsonValueSelectors`]). json_value_selectors: RefCell>, + /// The operands that appear in a query position, so the proxy can project + /// their payloads to query operands ([`crate::QueryOperands`]). Recorded + /// here rather than derived later because it is a fact about the statement's + /// shape, and the proxy needs it before it encrypts anything. + query_operands: RefCell>, + _ast: PhantomData<&'ast ()>, } @@ -74,6 +81,7 @@ impl<'ast> TypeInferencer<'ast> { scope_tracker: scope.into(), unifier: unifier.into(), json_value_selectors: RefCell::new(JsonValueSelectors::default()), + query_operands: RefCell::new(QueryOperands::default()), _ast: PhantomData, } } @@ -84,6 +92,19 @@ impl<'ast> TypeInferencer<'ast> { std::mem::take(&mut self.json_value_selectors.borrow_mut()) } + /// Takes the recorded query operands, leaving the inferencer's set empty. + pub(crate) fn take_query_operands(&self) -> QueryOperands<'ast> { + std::mem::take(&mut self.query_operands.borrow_mut()) + } + + pub(crate) fn record_query_operand_param(&self, param: Param) { + self.query_operands.borrow_mut().record_param(param); + } + + pub(crate) fn record_query_operand_literal(&self, node: &'ast sqltk::parser::ast::Value) { + self.query_operands.borrow_mut().record_literal(node); + } + pub(crate) fn record_json_value_selector_param( &self, param: Param, diff --git a/packages/eql-mapper/src/lib.rs b/packages/eql-mapper/src/lib.rs index a4f6e071f..b7666a2a6 100644 --- a/packages/eql-mapper/src/lib.rs +++ b/packages/eql-mapper/src/lib.rs @@ -10,6 +10,7 @@ mod json_value_selector; mod model; mod param; mod param_plan; +mod query_operands; mod renumber_params; mod scope_tracker; mod transformation_rules; @@ -24,6 +25,7 @@ pub use json_value_selector::*; pub use model::*; pub use param::*; pub use param_plan::*; +pub use query_operands::*; pub use type_checked_statement::*; pub use unifier::{ Array, AssociatedType, DomainIdentity, EqlTerm, EqlTermVariant, EqlTrait, EqlTraits, EqlValue, diff --git a/packages/eql-mapper/src/param_plan.rs b/packages/eql-mapper/src/param_plan.rs index 877d58a4d..e90202007 100644 --- a/packages/eql-mapper/src/param_plan.rs +++ b/packages/eql-mapper/src/param_plan.rs @@ -60,6 +60,11 @@ pub struct OutputParam { /// The input params it is built from. pub source: OutputParamSource, + + /// Whether this param is a **query operand** — an operand of a predicate, + /// which must reach PostgreSQL carrying only search terms and no + /// ciphertext. See [`crate::QueryOperands`]. + pub query_operand: bool, } /// The params of the rewritten statement, in the order PostgreSQL will see them. diff --git a/packages/eql-mapper/src/query_operands.rs b/packages/eql-mapper/src/query_operands.rs new file mode 100644 index 000000000..f167e83ac --- /dev/null +++ b/packages/eql-mapper/src/query_operands.rs @@ -0,0 +1,61 @@ +//! Which encrypted operands are **query operands** rather than stored values. +//! +//! An encrypted value reaches PostgreSQL in one of two shapes. A *stored value* +//! carries the column's whole payload — the ciphertext `c` plus every search +//! term. A *query operand* carries the terms but **never** a decryptable +//! ciphertext; the `eql_v3.query_*` domains enforce that with a `NOT (VALUE ? +//! 'c')` check. +//! +//! The two are indistinguishable by type: the operand of `WHERE col = $1` and +//! the value of `INSERT ... VALUES ($1)` are both [`crate::EqlTerm::Full`]. What +//! separates them is where they appear, which is what this records. +//! +//! The proxy needs it because it cannot encrypt a multi-term operand directly — +//! a single `EqlOperation::Query` yields only one term, so an operand is +//! encrypted in Store mode and then projected +//! (`EqlCiphertextV3::into_query_operand`). Without knowing the role it would +//! send a stored payload into a query position and PostgreSQL would reject it. + +use std::collections::HashSet; + +use sqltk::parser::ast; +use sqltk::NodeKey; + +use crate::Param; + +/// The set of operands that appear in a query position. +/// +/// Membership is decided syntactically, by the predicate an operand belongs to +/// — the same contexts whose rewrite rules cast to a `eql_v3.query_*` twin: +/// comparisons (`=`, `<>`, `<`, `<=`, `>`, `>=`), `LIKE`/`ILIKE` and `@@`. +/// Everything else — `INSERT` values, `UPDATE` assignments, containment needles +/// — is a stored value and keeps its full payload. +#[derive(Debug, Default)] +pub struct QueryOperands<'ast> { + params: HashSet, + literals: HashSet>, +} + +impl<'ast> QueryOperands<'ast> { + pub(crate) fn record_param(&mut self, param: Param) { + self.params.insert(param); + } + + pub(crate) fn record_literal(&mut self, node: &'ast ast::Value) { + self.literals.insert(NodeKey::new(node)); + } + + /// Whether the value bound to `param` is a query operand. + pub fn contains_param(&self, param: Param) -> bool { + self.params.contains(¶m) + } + + /// Whether the literal at `node` is a query operand. + pub fn contains_literal(&self, node: &'ast ast::Value) -> bool { + self.literals.contains(&NodeKey::new(node)) + } + + pub fn is_empty(&self) -> bool { + self.params.is_empty() && self.literals.is_empty() + } +} diff --git a/packages/eql-mapper/src/type_checked_statement.rs b/packages/eql-mapper/src/type_checked_statement.rs index 0209a5f64..9bcc37af9 100644 --- a/packages/eql-mapper/src/type_checked_statement.rs +++ b/packages/eql-mapper/src/type_checked_statement.rs @@ -4,6 +4,7 @@ use sqltk::parser::ast::{self, Statement}; use sqltk::{AsNodeKey, NodeKey, Transformable}; use crate::unifier::EqlTerm; +use crate::QueryOperands; use crate::{ CastFullPayloadOperands, DryRunnable, EqlMapperError, FailOnPlaceholderChange, JsonValueSelectors, OutputParam, OutputParamSource, Param, ParamPlan, PreserveEffectiveAliases, @@ -66,6 +67,14 @@ pub struct TypeCheckedStatement<'ast> { /// the proxy binds against. pub json_value_selectors: JsonValueSelectors<'ast>, + /// The operands that appear in a query position rather than a storing one. + /// + /// A query operand carries only search terms; a stored value carries the + /// ciphertext too. The two are the same [`EqlTerm`], so this is the only + /// thing telling the proxy which payload shape to send — see + /// [`QueryOperands`]. + pub query_operands: QueryOperands<'ast>, + /// A [`HashMap`] of AST node (using [`NodeKey`] as the key) to [`Type`]. The map contains a `Type` for every node /// in the AST with the node type is one of: [`Statement`], [`Query`], [`Insert`], [`Delete`], [`Expr`], /// [`SetExpr`], [`Select`], [`SelectItem`], [`Vec`], [`Function`], [`Values`], [`Value`]. @@ -91,6 +100,7 @@ impl<'ast> TypeCheckedStatement<'ast> { params: Vec<(Param, Value)>, literals: Vec<(EqlTerm, &'ast ast::Value)>, json_value_selectors: JsonValueSelectors<'ast>, + query_operands: QueryOperands<'ast>, node_types: Arc, Type>>, ) -> Self { Self { @@ -99,6 +109,7 @@ impl<'ast> TypeCheckedStatement<'ast> { params, literals, json_value_selectors, + query_operands, node_types, } } @@ -182,6 +193,7 @@ impl<'ast> TypeCheckedStatement<'ast> { param: Param((idx + 1) as u16), value, source, + query_operand: self.query_operands.contains_param(input), }) }) .collect::, EqlMapperError>>()?; From 477729e8b741f8a16327550fcb8e9958fcdb7a58 Mon Sep 17 00:00:00 2001 From: James Sadler Date: Sun, 26 Jul 2026 14:13:19 +1000 Subject: [PATCH 08/10] fix(mapper): GROUP BY an encrypted column must group by its equality term MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The third instance of the same trap. An encrypted column is a domain over `jsonb`, so a bare `GROUP BY col` grouped on the whole payload — including `c`, the ciphertext, which is randomised per encryption. Rows holding the same plaintext landed in different groups, degrading `GROUP BY` into `GROUP BY `: `select::group_by` saw 20 groups for 10 distinct values. Grouping is equality, so the key is the same term `=` uses — `eql_v3.eq_term(col)`, or `ord_term` for a domain that stores no `hm` (`eql_v3_integer_ord` groups through its ordering term, exactly as it compares). Once the key is `eq_term(col)`, PostgreSQL no longer sees the bare column as functionally dependent on it and rejects `SELECT col`. So a projection of a grouped column is lifted through an aggregate, as James specified: SELECT col, COUNT(*) FROM t GROUP BY col -> SELECT any_value(col) AS col, COUNT(*) FROM t GROUP BY eql_v3.eq_term(col) Every row in a group holds the same plaintext, so any of them decrypts to the right answer. `FIRST` does not exist in PostgreSQL and `min`/`max` are not defined over `jsonb`; `any_value` is polymorphic and works on the domain directly. It requires PostgreSQL 16+, but only for the projection case — grouping without selecting the column works on any version. The projection is matched on the resolved column rather than on syntax, so `SELECT t.col ... GROUP BY col` — which PostgreSQL accepts today — keeps working. The original projection name is preserved, so clients selecting by name are unaffected: this is what keeps `select::regression::select_with_order_by_in_subquery` (a `SELECT encrypted_text ... GROUP BY encrypted_text`) passing, which a group-key-only rewrite would have broken. Grouping by a column whose domain carries no equality term is now a capability error rather than a grouping on ciphertext. Integration 317 -> 323 passing, 87 -> 81 failing, nothing newly failing. Mapper units 107 -> 109. Stable-Commit-Id: q-4080dw66nm5aqm7g0q717mqvmw --- packages/eql-mapper/src/lib.rs | 87 +++++++++ .../src/transformation_rules/helpers.rs | 19 ++ .../src/transformation_rules/mod.rs | 2 + .../preserve_effective_aliases.rs | 39 ++-- .../rewrite_eql_group_by.rs | 183 ++++++++++++++++++ .../eql-mapper/src/type_checked_statement.rs | 7 +- 6 files changed, 319 insertions(+), 18 deletions(-) create mode 100644 packages/eql-mapper/src/transformation_rules/rewrite_eql_group_by.rs diff --git a/packages/eql-mapper/src/lib.rs b/packages/eql-mapper/src/lib.rs index b7666a2a6..20f17f522 100644 --- a/packages/eql-mapper/src/lib.rs +++ b/packages/eql-mapper/src/lib.rs @@ -2581,6 +2581,93 @@ mod test { } } + /// `GROUP BY` on an encrypted column groups by its equality term. A bare + /// `GROUP BY col` groups on the jsonb payload, whose ciphertext is + /// randomised per row, so equal plaintexts land in different groups. + /// + /// Projecting the grouped column has to be lifted through `any_value`, + /// because PostgreSQL no longer sees it as functionally dependent on the + /// group key — and the projection keeps the name the client asked for. + #[test] + fn group_by_encrypted_column_uses_eq_term() { + let schema = resolver(schema! { + tables: { + employees: { + id, + email (EQL("eql_v3_text_search"): Eq + Ord + TokenMatch), + salary (EQL("eql_v3_integer_ord"): Ord), + } + } + }); + + for (input, expected) in [ + // The column is not projected: only the group key changes. + ( + "SELECT COUNT(*) FROM employees GROUP BY email", + "SELECT COUNT(*) FROM employees GROUP BY eql_v3.eq_term(email)", + ), + // Projected: lifted through `any_value`, keeping its name. + ( + "SELECT email FROM employees GROUP BY email", + "SELECT any_value(email) AS email FROM employees GROUP BY eql_v3.eq_term(email)", + ), + // An explicit alias is preserved as-is. + ( + "SELECT email AS e FROM employees GROUP BY email", + "SELECT any_value(email) AS e FROM employees GROUP BY eql_v3.eq_term(email)", + ), + // Qualified projection of the same column still matches — the match + // is on the resolved column, not on syntax. + ( + "SELECT employees.email FROM employees GROUP BY email", + "SELECT any_value(employees.email) AS email FROM employees GROUP BY eql_v3.eq_term(email)", + ), + // A domain that stores no `hm` groups by its ordering term, the same + // fallback `=` uses. + ( + "SELECT COUNT(*) FROM employees GROUP BY salary", + "SELECT COUNT(*) FROM employees GROUP BY eql_v3.ord_term(salary)", + ), + // A native group key is left alone. + ( + "SELECT COUNT(*) FROM employees GROUP BY id", + "SELECT COUNT(*) FROM employees GROUP BY id", + ), + ] { + let statement = parse(input); + let typed = type_check(schema.clone(), &statement).unwrap(); + + assert_eq!( + typed.transform(HashMap::new()).unwrap().to_string(), + expected, + "unexpected rewrite for `{input}`" + ); + } + } + + /// Grouping by a column whose domain carries no equality term is a + /// capability error, not a grouping on ciphertext. + #[test] + fn group_by_column_without_equality_term_is_an_error() { + let schema = resolver(schema! { + tables: { + employees: { + id, + name (EQL("eql_v3_text_match"): TokenMatch), + } + } + }); + + let statement = parse("SELECT COUNT(*) FROM employees GROUP BY name"); + let typed = type_check(schema, &statement).unwrap(); + + let err = typed.transform(HashMap::new()).unwrap_err().to_string(); + assert!( + err.contains("GROUP BY") && err.contains("no equality term"), + "expected a capability error, got: {err}" + ); + } + /// A block-ORE domain orders through `ord_term_ore`. #[test] fn order_by_ore_column_uses_ord_term_ore() { diff --git a/packages/eql-mapper/src/transformation_rules/helpers.rs b/packages/eql-mapper/src/transformation_rules/helpers.rs index d6b8f8dda..db8273f75 100644 --- a/packages/eql-mapper/src/transformation_rules/helpers.rs +++ b/packages/eql-mapper/src/transformation_rules/helpers.rs @@ -138,6 +138,25 @@ pub(crate) fn cast_expr_to_v3_domain(wrapped: Expr, schema: &str, domain: &str) } } +/// Builds `()` — an unqualified, single-argument function call +/// resolved from `search_path` (i.e. a built-in such as `any_value`). +pub(crate) fn pg_catalog_fn_call(fn_name: &str, arg: Expr) -> Expr { + Expr::Function(Function { + name: ObjectName(vec![ObjectNamePart::Identifier(Ident::new(fn_name))]), + uses_odbc_syntax: false, + args: FunctionArguments::List(FunctionArgumentList { + args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(arg))], + duplicate_treatment: None, + clauses: vec![], + }), + parameters: FunctionArguments::None, + filter: None, + null_treatment: None, + over: None, + within_group: vec![], + }) +} + /// Builds `eql_v3.()` — a call to an EQL v3 term-extraction function /// (`eq_term`, `ord_term`, `ord_term_ore`, `match_term`). pub(crate) fn eql_v3_term_call(fn_name: &str, arg: Expr) -> Expr { diff --git a/packages/eql-mapper/src/transformation_rules/mod.rs b/packages/eql-mapper/src/transformation_rules/mod.rs index f3ff77fcc..6f1a412ee 100644 --- a/packages/eql-mapper/src/transformation_rules/mod.rs +++ b/packages/eql-mapper/src/transformation_rules/mod.rs @@ -16,6 +16,7 @@ mod fail_on_placeholder_change; mod preserve_effective_aliases; mod rewrite_containment_ops; mod rewrite_eql_comparison_ops; +mod rewrite_eql_group_by; mod rewrite_eql_match_ops; mod rewrite_eql_order_by; mod rewrite_json_value_selector_eq; @@ -29,6 +30,7 @@ pub(crate) use fail_on_placeholder_change::*; pub(crate) use preserve_effective_aliases::*; pub(crate) use rewrite_containment_ops::*; pub(crate) use rewrite_eql_comparison_ops::*; +pub(crate) use rewrite_eql_group_by::*; pub(crate) use rewrite_eql_match_ops::*; pub(crate) use rewrite_eql_order_by::*; pub(crate) use rewrite_json_value_selector_eq::*; diff --git a/packages/eql-mapper/src/transformation_rules/preserve_effective_aliases.rs b/packages/eql-mapper/src/transformation_rules/preserve_effective_aliases.rs index b73d5cae0..fd8b39ad3 100644 --- a/packages/eql-mapper/src/transformation_rules/preserve_effective_aliases.rs +++ b/packages/eql-mapper/src/transformation_rules/preserve_effective_aliases.rs @@ -126,23 +126,32 @@ impl PreserveEffectiveAliases { } fn derive_effective_alias(node: &SelectItem) -> Option { - match node { - SelectItem::UnnamedExpr(expr) => Self::derive_effective_alias_for_expr(expr), - SelectItem::ExprWithAlias { expr: _, alias } => Some(alias.clone()), - _ => None, - } + derive_effective_alias(node) } +} - fn derive_effective_alias_for_expr(expr: &Expr) -> Option { - match expr { - Expr::Identifier(ident) => Some(ident.clone()), - Expr::CompoundIdentifier(idents) => Some(idents.last().unwrap().clone()), - Expr::Function(Function { name, .. }) => { - let ObjectNamePart::Identifier(ident) = name.0.last().unwrap().clone(); - Some(ident) - } - Expr::Nested(expr) => Self::derive_effective_alias_for_expr(expr), - _ => None, +/// The name PostgreSQL would give a projection column — its explicit alias, or +/// the name derived from the expression. +/// +/// Shared with [`super::RewriteEqlGroupBy`], which wraps a projected column in +/// an aggregate and must give the result the name the client asked for. +pub(crate) fn derive_effective_alias(node: &SelectItem) -> Option { + match node { + SelectItem::UnnamedExpr(expr) => derive_effective_alias_for_expr(expr), + SelectItem::ExprWithAlias { expr: _, alias } => Some(alias.clone()), + _ => None, + } +} + +fn derive_effective_alias_for_expr(expr: &Expr) -> Option { + match expr { + Expr::Identifier(ident) => Some(ident.clone()), + Expr::CompoundIdentifier(idents) => Some(idents.last().unwrap().clone()), + Expr::Function(Function { name, .. }) => { + let ObjectNamePart::Identifier(ident) = name.0.last().unwrap().clone(); + Some(ident) } + Expr::Nested(expr) => derive_effective_alias_for_expr(expr), + _ => None, } } diff --git a/packages/eql-mapper/src/transformation_rules/rewrite_eql_group_by.rs b/packages/eql-mapper/src/transformation_rules/rewrite_eql_group_by.rs new file mode 100644 index 000000000..4c5466c9e --- /dev/null +++ b/packages/eql-mapper/src/transformation_rules/rewrite_eql_group_by.rs @@ -0,0 +1,183 @@ +use std::collections::HashMap; +use std::mem; +use std::sync::Arc; + +use sqltk::parser::ast::Value as SqltkValue; +use sqltk::parser::ast::{Expr, GroupByExpr, Select, SelectItem, ValueWithSpan}; +use sqltk::parser::tokenizer::Span; +use sqltk::{NodeKey, NodePath, Visitable}; + +use crate::unifier::{EqlValue, Type, Value}; +use crate::EqlMapperError; + +use super::helpers::{eql_v3_term_call, pg_catalog_fn_call}; +use super::preserve_effective_aliases::derive_effective_alias; +use super::TransformationRule; + +/// Rewrites `GROUP BY` on an encrypted column to group by its **equality term**, +/// and lifts any projection of that column through an aggregate so the query +/// stays valid: +/// +/// ```sql +/// SELECT col, COUNT(*) FROM t GROUP BY col +/// -- becomes +/// SELECT any_value(col) AS col, COUNT(*) FROM t GROUP BY eql_v3.eq_term(col) +/// ``` +/// +/// **Without this rewrite the grouping is silently wrong.** An encrypted column +/// is a domain over `jsonb`, so a bare `GROUP BY` groups on the whole payload — +/// including `c`, the ciphertext, which is randomised per encryption. Two rows +/// holding the *same* plaintext land in different groups, so `GROUP BY` degrades +/// into `GROUP BY `. +/// +/// Grouping is equality, so the key is the same `eq_term` an `=` comparison uses +/// (`ord_term` for a domain that stores no `hm`). +/// +/// Once the key is `eq_term(col)`, PostgreSQL no longer sees the bare column as +/// functionally dependent on it and would reject `SELECT col`. `any_value` +/// resolves that: every row in a group holds the same plaintext, so any of them +/// decrypts to the right answer. The original projection name is preserved, so +/// clients selecting by name are unaffected. +/// +/// `any_value` requires PostgreSQL 16 or later. Only the projection case needs +/// it; grouping without selecting the column works on any version. +#[derive(Debug)] +pub struct RewriteEqlGroupBy<'ast> { + node_types: Arc, Type>>, +} + +impl<'ast> RewriteEqlGroupBy<'ast> { + pub fn new(node_types: Arc, Type>>) -> Self { + Self { node_types } + } + + fn eql_value_of(&self, expr: &'ast Expr) -> Option { + match self.node_types.get(&NodeKey::new(expr)) { + Some(Type::Value(Value::Eql(eql_term))) => Some(eql_term.eql_value().clone()), + _ => None, + } + } + + /// The encrypted columns a `GROUP BY` groups on, in order. + fn grouped_eql_values(&self, group_by: &'ast GroupByExpr) -> Vec> { + match group_by { + GroupByExpr::Expressions(exprs, _) => { + exprs.iter().map(|expr| self.eql_value_of(expr)).collect() + } + GroupByExpr::All(_) => vec![], + } + } + + /// The expression a select item projects, if it is a plain one. + fn select_item_expr(item: &'ast SelectItem) -> Option<&'ast Expr> { + match item { + SelectItem::UnnamedExpr(expr) | SelectItem::ExprWithAlias { expr, .. } => Some(expr), + _ => None, + } + } + + /// Whether `item` projects one of the grouped encrypted columns, and so must + /// be lifted through an aggregate. + /// + /// Matched on the resolved column rather than on syntax, so `SELECT t.col … + /// GROUP BY col` — which PostgreSQL accepts today — keeps working. + fn projects_grouped_column(&self, item: &'ast SelectItem, grouped: &[EqlValue]) -> bool { + Self::select_item_expr(item) + .and_then(|expr| self.eql_value_of(expr)) + .is_some_and(|value| grouped.contains(&value)) + } +} + +impl<'ast> TransformationRule<'ast> for RewriteEqlGroupBy<'ast> { + fn apply( + &mut self, + node_path: &NodePath<'ast>, + target_node: &mut N, + ) -> Result { + // Read the encrypted columns from the ORIGINAL select — `node_types` is + // keyed by it, and the target's children are already rewritten. + let Some((original,)) = node_path.last_1_as::() else { + return Ok(false); + }; + + // Group by the equality term. + if let GroupByExpr::Expressions(exprs, _) = &mut target.group_by { + for (expr, eql_value) in exprs.iter_mut().zip(grouped.iter()) { + let Some(eql_value) = eql_value else { continue }; + + let identity = eql_value.domain_identity(); + let Some(term_fn) = identity.eq_term_fn() else { + return Err(EqlMapperError::Transform(format!( + "encrypted column {} cannot be used in GROUP BY (domain {} carries no equality term)", + identity.token, identity.domain.value + ))); + }; + + let grouped_expr = mem::replace( + expr, + Expr::Value(ValueWithSpan { + value: SqltkValue::Null, + span: Span::empty(), + }), + ); + *expr = eql_v3_term_call(term_fn, grouped_expr); + } + } + + // Lift any projection of a grouped column through `any_value`, keeping + // the name the client asked for. + let grouped: Vec = grouped.into_iter().flatten().collect(); + for (original_item, target_item) in + original.projection.iter().zip(target.projection.iter_mut()) + { + if !self.projects_grouped_column(original_item, &grouped) { + continue; + } + + let alias = derive_effective_alias(original_item); + let (SelectItem::UnnamedExpr(expr) | SelectItem::ExprWithAlias { expr, .. }) = + target_item + else { + continue; + }; + + let projected = mem::replace( + expr, + Expr::Value(ValueWithSpan { + value: SqltkValue::Null, + span: Span::empty(), + }), + ); + let aggregated = pg_catalog_fn_call("any_value", projected); + + *target_item = match alias { + Some(alias) => SelectItem::ExprWithAlias { + expr: aggregated, + alias, + }, + None => SelectItem::UnnamedExpr(aggregated), + }; + } + + Ok(true) + } + + fn would_edit(&mut self, node_path: &NodePath<'ast>, _target_node: &N) -> bool { + match node_path.last_1_as::