From ef6f7d2e079b5359e5341af90eee3f6850c06411 Mon Sep 17 00:00:00 2001 From: Lee Hannigan Date: Fri, 18 Sep 2026 11:57:38 +0000 Subject: [PATCH] fix(storage): store U+0000 in strings on PostgreSQL and MongoDB DynamoDB accepts the character U+0000 anywhere a string appears: keys, index keys, values, strings in lists, map keys, attribute names (measured against the service 2026-09-18; it sorts as the byte 0x00 and is a different key from the six-character text "\u0000"). PostgreSQL TEXT rejects the byte and jsonb rejects the escape, so the backend answered 500 for every such item and for a Query whose key value carried it. BSON field names are C strings, so MongoDB answered 500 when the character appeared in an attribute name or map key. Shared escape (crates/storage/src/util/control_chars.rs): U+0000 is stored as U+0001 U+0001, U+0001 as U+0001 U+0002, everything else unchanged. Exhaustive tests over every string up to length 4 on {U+0000, U+0001, U+0002, a} pin round trip, UTF-8 byte order, prefix preservation, and distinctness, so COLLATE "C" comparisons, BETWEEN, begins_with, and the row-comparison page cursors return the rows the raw text would. PostgreSQL: item documents go through item_to_json/json_to_item (every item_data write, stream records, gsi_pending rows and contexts, vector rows); key column text goes through data/key_text.rs wrappers at every producing site; the TTL index and sweep and the vector search filters address attributes by the escaped name. Migration 004_escape_control_chars re-encodes rows written by earlier releases that contain U+0001, one table per transaction with a per-table progress marker, only rows containing U+0001 touched, predicate independent of standard_conforming_strings. The server refuses to start on a data database that does not record the migration (rows this build writes are indistinguishable from legacy rows to the migration), extenddb init records it on fresh deployments, and extenddb migrate refuses to run it while other sessions hold the data database. MongoDB: item_to_document and its inverse escape object keys only (BSON string values already hold the byte); conditions and native updates on names with these characters use the existing in-process fallback paths, as names with "." or "$" already do; field paths with no fallback use the escaped name. Tests: tests/test_nul_strings.py (14 wire cases passing against DynamoDB and all three backends, plus a stream record case on ExtendDB servers); unit tests on the escape, the key wrappers, the MongoDB document conversion and predicates; live migration tests for re-encoding, the second-run skip, the standard_conforming_strings setting, the startup gate, and the connection guard. Docs: storage design, upgrade manual. Closes #364 --- crates/storage-mongodb/src/data/mod.rs | 137 ++- crates/storage-mongodb/src/data_engine.rs | 38 +- crates/storage-mongodb/src/metadata_engine.rs | 25 +- crates/storage-mongodb/src/pushdown.rs | 44 +- crates/storage-mongodb/src/stream_engine.rs | 5 +- crates/storage-postgres/src/bootstrapper.rs | 3 + .../storage-postgres/src/data/delete_item.rs | 4 +- crates/storage-postgres/src/data/index.rs | 9 +- crates/storage-postgres/src/data/key_text.rs | 139 +++ crates/storage-postgres/src/data/mod.rs | 40 +- crates/storage-postgres/src/data/put_item.rs | 9 +- crates/storage-postgres/src/data/query.rs | 4 +- .../storage-postgres/src/data/query_scan.rs | 6 +- .../storage-postgres/src/data/tx_helpers.rs | 15 +- .../storage-postgres/src/data/update_item.rs | 8 +- .../storage-postgres/src/data/vector_index.rs | 11 +- crates/storage-postgres/src/gsi_queue.rs | 28 +- crates/storage-postgres/src/lib.rs | 44 + .../storage-postgres/src/metadata_engine.rs | 8 +- crates/storage-postgres/src/migrations.rs | 1037 ++++++++++++++++- crates/storage-postgres/src/stream_engine.rs | 7 +- crates/storage-postgres/src/vector_search.rs | 15 +- crates/storage/src/util/control_chars.rs | 332 ++++++ crates/storage/src/util/key.rs | 9 +- crates/storage/src/util/mod.rs | 5 + docs/design/04-component-storage.md | 19 + docs/manuals/07-upgrade-manual.md | 60 + tests/test_nul_strings.py | 378 ++++++ 28 files changed, 2343 insertions(+), 96 deletions(-) create mode 100644 crates/storage-postgres/src/data/key_text.rs create mode 100644 crates/storage/src/util/control_chars.rs create mode 100644 tests/test_nul_strings.py diff --git a/crates/storage-mongodb/src/data/mod.rs b/crates/storage-mongodb/src/data/mod.rs index 45465ba61..2036586ce 100644 --- a/crates/storage-mongodb/src/data/mod.rs +++ b/crates/storage-mongodb/src/data/mod.rs @@ -13,7 +13,9 @@ use extenddb_core::types::{ AttributeDefinition, AttributeValue, Item, KeySchemaElement, ScalarAttributeType, }; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{composite_pk_to_text, encode_netstring_composite, sk_info}; +use extenddb_storage::util::{ + composite_pk_to_text, encode_netstring_composite, escape_json_keys, sk_info, unescape_json_keys, +}; /// Returns the `MongoDB` collection name for a `DynamoDB` table. pub fn data_collection_name(table_id: &str) -> String { @@ -32,6 +34,42 @@ pub fn composite_id(pk_text: &str, sk_text: &str) -> String { encode_netstring_composite(&[pk_text.to_owned(), sk_text.to_owned()]) } +/// Convert an item (base item or projected index item) to the BSON stored +/// under `item_data`. +/// +/// BSON field names are C strings and cannot contain U+0000, while `DynamoDB` +/// accepts it in attribute names and map keys. Every object key in the item's +/// JSON form is run through the shared order-preserving escape (U+0000 and +/// U+0001 become two-character sequences, see +/// `extenddb_storage::util::escape_control`) before the BSON conversion, so +/// the stored document never carries a NUL field name. String values are left +/// alone; BSON strings are length-prefixed and hold U+0000 verbatim. +/// +/// Every conversion of an item to `item_data` BSON must go through this +/// function or [`json_to_item_data`]; a direct `bson::to_bson` on item JSON +/// fails in the driver at write time when a name contains U+0000. +/// [`item_data_to_json`] reverses the escape on read. +pub(crate) fn item_to_item_data(item: &Item) -> Result { + let item_json = + serde_json::to_value(item).map_err(|e| StorageError::Internal(e.to_string()))?; + json_to_item_data(item_json) +} + +/// JSON-tree variant of [`item_to_item_data`], for payloads that carry items +/// nested inside another structure (stream records with key/image maps). +pub(crate) fn json_to_item_data(json: serde_json::Value) -> Result { + let escaped = escape_json_keys(json); + bson::to_bson(&escaped).map_err(|e| StorageError::Internal(e.to_string())) +} + +/// Reverse of [`json_to_item_data`]: BSON read from storage back to JSON with +/// the object-key escape undone. +pub(crate) fn item_data_to_json(value: &bson::Bson) -> Result { + let json: serde_json::Value = bson::from_bson(value.clone()) + .map_err(|e| StorageError::Internal(format!("BSON to JSON conversion error: {e}")))?; + Ok(unescape_json_keys(json)) +} + /// Convert a `DynamoDB` Item to a `MongoDB` BSON document for storage. /// /// Document structure: `{ _id, pk, sk_s/sk_n/sk_b, item_data }` @@ -42,10 +80,9 @@ pub fn item_to_document( ) -> Result { let pk_text = composite_pk_to_text(item, key_schema)?; - // Serialize the full item as item_data - let item_json = - serde_json::to_value(item).map_err(|e| StorageError::Internal(e.to_string()))?; - let item_bson = bson::to_bson(&item_json).map_err(|e| StorageError::Internal(e.to_string()))?; + // Serialize the full item as item_data (object keys escaped, see + // item_to_item_data). + let item_bson = item_to_item_data(item)?; let mut doc = Document::new(); @@ -192,9 +229,7 @@ pub fn index_document( insert_typed_sk(&mut doc, &field, sk_type, sk_value)?; } - let item_json = - serde_json::to_value(projected).map_err(|e| StorageError::Internal(e.to_string()))?; - let item_bson = bson::to_bson(&item_json).map_err(|e| StorageError::Internal(e.to_string()))?; + let item_bson = item_to_item_data(projected)?; doc.insert("item_data", item_bson); Ok(doc) @@ -254,8 +289,7 @@ pub fn document_to_item(doc: &Document) -> Result { .get("item_data") .ok_or_else(|| StorageError::Internal("Document missing item_data field".to_string()))?; - let json_value: serde_json::Value = bson::from_bson(item_data.clone()) - .map_err(|e| StorageError::Internal(format!("BSON to JSON conversion error: {e}")))?; + let json_value = item_data_to_json(item_data)?; let item: Item = serde_json::from_value(json_value) .map_err(|e| StorageError::Internal(format!("JSON to Item conversion error: {e}")))?; @@ -434,6 +468,89 @@ mod tests { assert!(matches!(err, StorageError::Validation(_))); } + /// Walk a BSON tree and assert no document field name contains U+0000. + fn assert_no_nul_field_names(value: &bson::Bson) { + match value { + bson::Bson::Document(doc) => { + for (k, v) in doc { + assert!(!k.contains('\u{0}'), "field name contains NUL: {k:?}"); + assert_no_nul_field_names(v); + } + } + bson::Bson::Array(items) => { + for item in items { + assert_no_nul_field_names(item); + } + } + _ => {} + } + } + + #[test] + fn item_document_round_trips_control_characters_in_names() { + use std::collections::BTreeMap; + + let (schema, attrs) = schema_pk_str_sk_num(); + + let mut inner = BTreeMap::new(); + inner.insert( + "deep\u{0}key".to_owned(), + AttributeValue::S("v\u{0}".to_owned()), + ); + let mut mid = BTreeMap::new(); + mid.insert("mid\u{1}key".to_owned(), AttributeValue::M(inner)); + let mut in_list = BTreeMap::new(); + in_list.insert( + "list\u{0}\u{1}key".to_owned(), + AttributeValue::S("w".to_owned()), + ); + + let mut item = Item::new(); + item.insert("pk".to_owned(), AttributeValue::S("p".to_owned())); + item.insert("sk".to_owned(), AttributeValue::N("1".to_owned())); + item.insert( + "top\u{0}\u{1}name".to_owned(), + AttributeValue::S("x".to_owned()), + ); + item.insert("map".to_owned(), AttributeValue::M(mid)); + item.insert( + "list".to_owned(), + AttributeValue::L(vec![AttributeValue::M(in_list)]), + ); + + let doc = item_to_document(&item, &schema, &attrs).unwrap(); + assert_no_nul_field_names(&bson::Bson::Document(doc.clone())); + + let back = document_to_item(&doc).unwrap(); + assert_eq!(back, item); + } + + #[test] + fn plain_item_document_stores_raw_field_names() { + use std::collections::BTreeMap; + + let (schema, attrs) = schema_pk_str_sk_num(); + + let mut inner = BTreeMap::new(); + inner.insert("k".to_owned(), AttributeValue::S("v".to_owned())); + let mut item = Item::new(); + item.insert("pk".to_owned(), AttributeValue::S("p".to_owned())); + item.insert("sk".to_owned(), AttributeValue::N("1".to_owned())); + item.insert("attr".to_owned(), AttributeValue::M(inner)); + + let doc = item_to_document(&item, &schema, &attrs).unwrap(); + let item_data = doc.get_document("item_data").unwrap(); + let names: Vec<&str> = item_data.keys().map(String::as_str).collect(); + assert_eq!(names, vec!["attr", "pk", "sk"]); + let attr_m = item_data + .get_document("attr") + .unwrap() + .get_document("M") + .unwrap(); + let map_keys: Vec<&str> = attr_m.keys().map(String::as_str).collect(); + assert_eq!(map_keys, vec!["k"]); + } + #[test] fn composite_id_disambiguates_delimiter_in_pk_or_sk() { // Two items whose naive "{pk}#{sk}" strings would collide must diff --git a/crates/storage-mongodb/src/data_engine.rs b/crates/storage-mongodb/src/data_engine.rs index 3895c4562..af40046fd 100644 --- a/crates/storage-mongodb/src/data_engine.rs +++ b/crates/storage-mongodb/src/data_engine.rs @@ -18,7 +18,7 @@ use extenddb_core::types::{ }; use extenddb_storage::error::StorageError; use extenddb_storage::util::{ - composite_pk_to_text, encode_netstring_composite, pk_to_text, sk_info, + composite_pk_to_text, encode_netstring_composite, needs_escape, pk_to_text, sk_info, }; use extenddb_storage::{ DataEngine, IdempotencyKey, ItemPairResult, QueryResult, StreamCapture, TransactGetOp, @@ -29,7 +29,7 @@ use crate::MongoEngine; use crate::condition::condition_to_filter; use crate::data::{ binary_sk_to_hex, composite_id, data_collection_name, document_to_item, index_document, - index_entry_filter, item_to_document, pk_filter, sk_field_name, sk_suffix, + index_entry_filter, item_to_document, json_to_item_data, pk_filter, sk_field_name, sk_suffix, }; use crate::pushdown::{Pushable, is_pushable}; @@ -37,8 +37,9 @@ use extenddb_core::types::{Projection, ProjectionType}; /// Resolve a single-component update path for the native MongoDB fast path. /// Literal attribute names containing dots or beginning with `$` cannot be -/// represented safely by MongoDB's ordinary update paths, so they must use -/// the Rust/session path. +/// represented safely by MongoDB's ordinary update paths, and names +/// containing U+0000 or U+0001 differ from the stored (escaped) field name, +/// so they must use the Rust/session path. fn native_attribute_name(path: &[PathElement], maps: &ExpressionMaps) -> Option { if path.len() != 1 { return None; @@ -47,7 +48,8 @@ fn native_attribute_name(path: &[PathElement], maps: &ExpressionMaps) -> Option< return None; }; let attr_name = resolve_name_ref(raw_name, maps).ok()?; - (!attr_name.contains('.') && !attr_name.starts_with('$')).then(|| attr_name.into_owned()) + (!attr_name.contains('.') && !attr_name.starts_with('$') && !needs_escape(&attr_name)) + .then(|| attr_name.into_owned()) } #[derive(Clone, Copy)] @@ -1651,8 +1653,11 @@ impl MongoEngine { _ => return None, // complex expressions (if_not_exists, list_append, arithmetic) }; let field = format!("item_data.{attr_name}"); + // The value lands inside item_data, so its nested map + // keys must be in stored (escaped) form, matching what + // item_to_document writes. let val_json = serde_json::to_value(val).ok()?; - let val_bson = bson::to_bson(&val_json).ok()?; + let val_bson = json_to_item_data(val_json).ok()?; set_doc.insert(field, val_bson); } UpdateAction::Remove { path } => { @@ -2033,8 +2038,10 @@ impl MongoEngine { let record_json = serde_json::to_value(&record).map_err(|e| StorageError::Internal(e.to_string()))?; - let record_bson = - bson::to_bson(&record_json).map_err(|e| StorageError::Internal(e.to_string()))?; + // The record's key and image maps carry attribute names as object + // keys, which become BSON field names; convert through the same + // escaping helper as item_data. + let record_bson = json_to_item_data(record_json)?; // key_info already carries table_id — no need to re-read the catalog // just to resolve it, and re-reading inside the session against the @@ -4385,6 +4392,21 @@ transaction_read_concern = "majority""#, assert_eq!(native_attribute_name(&path, &maps), None); } + #[test] + fn native_update_rejects_control_character_literal_attribute_names() { + // Stored field names carry the escaped form of U+0000 and U+0001, + // so a native update path built from the raw name would target a + // different field; these names must use the Rust/session path. + for name in ["a\u{0}b", "a\u{1}b"] { + let mut names = std::collections::HashMap::new(); + names.insert("name".to_owned(), name.to_owned()); + let maps = ExpressionMaps::new(names, std::collections::HashMap::new()); + let path = vec![PathElement::Attribute("#name".to_owned())]; + + assert_eq!(native_attribute_name(&path, &maps), None, "{name:?}"); + } + } + #[test] fn native_update_keeps_simple_attribute_names_on_fast_path() { let maps = ExpressionMaps::new( diff --git a/crates/storage-mongodb/src/metadata_engine.rs b/crates/storage-mongodb/src/metadata_engine.rs index 7fe21f9c3..3dac4241b 100644 --- a/crates/storage-mongodb/src/metadata_engine.rs +++ b/crates/storage-mongodb/src/metadata_engine.rs @@ -18,6 +18,7 @@ use extenddb_storage::error::StorageError; use crate::MongoEngine; use crate::data::{data_collection_name, document_to_item}; +use extenddb_storage::util::escape_control; fn extract_id_fields(doc: &Document) -> (String, String) { let id = doc.get_document("_id").ok(); @@ -35,13 +36,15 @@ fn extract_id_fields(doc: &Document) -> (String, String) { /// Build the expression that reads a DynamoDB attribute from `item_data` as a /// literal field name. MongoDB normally interprets dots in a query field path /// as nested-document traversal, but DynamoDB permits dots in attribute names. +/// The stored field name is the escaped form of the attribute name (see +/// `extenddb_storage::util::escape_control`), so the lookup targets that form. fn literal_ttl_value_expression(ttl_attribute: &str) -> Document { doc! { "$getField": { "field": "N", "input": { "$getField": { - "field": ttl_attribute, + "field": escape_control(ttl_attribute).as_ref(), "input": "$item_data", } } @@ -69,7 +72,7 @@ fn literal_ttl_epoch_expression(ttl_attribute: &str) -> Document { /// also performs the expiry comparison in MongoDB, so non-expired low-`_id` /// documents cannot starve expired documents later in the collection. fn ttl_candidate_query(ttl_attribute: &str, now_epoch: i64) -> (Document, Document) { - let ttl_field = format!("item_data.{ttl_attribute}.N"); + let ttl_field = format!("item_data.{}.N", escape_control(ttl_attribute)); if ttl_attribute.contains('.') { let ttl_epoch = literal_ttl_epoch_expression(ttl_attribute); ( @@ -374,8 +377,12 @@ impl MetadataEngine for MongoEngine { let coll_name = data_collection_name(table_id); let data_coll = self.data_db.collection::(&coll_name); - let index_name = format!("idx_ttl_{ttl_attribute}"); - let index_key = format!("item_data.{ttl_attribute}.N"); + // The stored field name is the escaped attribute name, and the + // escape also keeps U+0000 out of the index name. Names without + // U+0000 or U+0001 escape to themselves, so existing index names + // do not change. + let index_name = format!("idx_ttl_{}", escape_control(&ttl_attribute)); + let index_key = format!("item_data.{}.N", escape_control(&ttl_attribute)); let index = IndexModel::builder() .keys(doc! { &index_key: 1 }) @@ -433,7 +440,9 @@ impl MetadataEngine for MongoEngine { .get_str("table_id") .map_err(|_| StorageError::Internal("missing table_id".to_string()))?; - let index_name = format!("idx_ttl_{ttl_attribute}"); + // Same escaped form as create_ttl_index so the drop targets the + // index that was created. + let index_name = format!("idx_ttl_{}", escape_control(&ttl_attribute)); let coll_name = data_collection_name(table_id); let data_coll = self.data_db.collection::(&coll_name); @@ -513,9 +522,11 @@ impl MetadataEngine for MongoEngine { if items.len() >= limit { break; } - // Parse the TTL value and check if expired + // Parse the TTL value and check if expired. The stored field + // name is the escaped attribute name. if let Ok(item_data) = doc.get_document("item_data") - && let Ok(ttl_obj) = item_data.get_document(&ttl_attribute) + && let Ok(ttl_obj) = + item_data.get_document(escape_control(&ttl_attribute).as_ref()) && let Ok(n_str) = ttl_obj.get_str("N") && let Ok(ttl_val) = n_str.parse::() && ttl_val >= 1 diff --git a/crates/storage-mongodb/src/pushdown.rs b/crates/storage-mongodb/src/pushdown.rs index 616ef7508..05f96dac9 100644 --- a/crates/storage-mongodb/src/pushdown.rs +++ b/crates/storage-mongodb/src/pushdown.rs @@ -44,6 +44,7 @@ use extenddb_core::expression::{CompareOp, Expr, ExpressionMaps, PathElement}; use extenddb_core::types::AttributeValue; +use extenddb_storage::util::needs_escape; /// Outcome of the pushdown analyzer. #[derive(Debug, Clone, PartialEq, Eq)] @@ -76,7 +77,10 @@ pub fn is_pushable(expr: &Expr, maps: &ExpressionMaps) -> Pushable { /// Return why an expression contains a literal DynamoDB attribute name that /// cannot safely be emitted as a MongoDB field path. A dot would mean nested /// document traversal in MongoDB, while a leading `$` would be interpreted as -/// an operator or otherwise have special meaning in a MongoDB path. +/// an operator or otherwise have special meaning in a MongoDB path. A name +/// containing U+0000 or U+0001 is stored under its escaped form (see +/// `extenddb_storage::util::escape_control`), so the raw name would not match +/// the stored field. fn unsafe_attribute_reason(expr: &Expr, maps: &ExpressionMaps) -> Option<&'static str> { match expr { Expr::Path(elements) => elements.iter().find_map(|element| match element { @@ -91,6 +95,9 @@ fn unsafe_attribute_reason(expr: &Expr, maps: &ExpressionMaps) -> Option<&'stati Some(name) if name.starts_with('$') => { Some("literal attribute name begins with '$'") } + Some(name) if needs_escape(name) => { + Some("literal attribute name contains U+0000 or U+0001") + } Some(_) => None, None => Some("attribute name alias is unresolved"), } @@ -393,6 +400,41 @@ mod tests { ); } + #[test] + fn nul_attribute_name_is_not_pushable() { + // The stored field name is the escaped form, so a raw path on a + // name containing U+0000 would never match; must fall back. + let expr = Expr::Compare { + left: Box::new(path("a\u{0}b")), + op: CompareOp::Eq, + right: Box::new(Expr::Placeholder(":value".into())), + }; + let maps = maps_with(&[(":value", AttributeValue::S("value".into()))]); + + assert_eq!( + is_pushable(&expr, &maps), + Pushable::No("literal attribute name contains U+0000 or U+0001") + ); + } + + #[test] + fn control_character_attribute_alias_is_not_pushable() { + for resolved in ["a\u{0}b", "a\u{1}b"] { + let expr = Expr::Compare { + left: Box::new(path("#name")), + op: CompareOp::Eq, + right: Box::new(Expr::Placeholder(":value".into())), + }; + let maps = aliased_maps(resolved, AttributeValue::S("value".into())); + + assert_eq!( + is_pushable(&expr, &maps), + Pushable::No("literal attribute name contains U+0000 or U+0001"), + "{resolved:?}" + ); + } + } + #[test] fn nested_document_path_without_literal_dot_remains_pushable() { let expr = Expr::Compare { diff --git a/crates/storage-mongodb/src/stream_engine.rs b/crates/storage-mongodb/src/stream_engine.rs index e976a667a..b57180c18 100644 --- a/crates/storage-mongodb/src/stream_engine.rs +++ b/crates/storage-mongodb/src/stream_engine.rs @@ -368,8 +368,9 @@ impl StreamEngine for MongoEngine { let record_bson = d .get("record_data") .ok_or_else(|| StorageError::Internal("Missing record_data".to_owned()))?; - let json_val: serde_json::Value = bson::from_bson(record_bson.clone()) - .map_err(|e| StorageError::Internal(e.to_string()))?; + // Reverses the object-key escape applied when the record + // was written; see json_to_item_data. + let json_val = crate::data::item_data_to_json(record_bson)?; serde_json::from_value(json_val) .map_err(|e| StorageError::Internal(e.to_string())) }) diff --git a/crates/storage-postgres/src/bootstrapper.rs b/crates/storage-postgres/src/bootstrapper.rs index a8827df35..e25ae3f02 100755 --- a/crates/storage-postgres/src/bootstrapper.rs +++ b/crates/storage-postgres/src/bootstrapper.rs @@ -110,6 +110,9 @@ impl PostgresBootstrapper { .username(&self.config.app_user) .password(&self.config.app_password) .database(database) + // Lets a migration tell its own connections apart from a server's + // in pg_stat_activity (see `migrations::MIGRATE_APPLICATION_NAME`). + .application_name(migrations::MIGRATE_APPLICATION_NAME) } /// Build the connection URL for the application user and a named database. diff --git a/crates/storage-postgres/src/data/delete_item.rs b/crates/storage-postgres/src/data/delete_item.rs index 1ea9dd1e8..e350194ac 100755 --- a/crates/storage-postgres/src/data/delete_item.rs +++ b/crates/storage-postgres/src/data/delete_item.rs @@ -7,7 +7,9 @@ use extenddb_core::expression::{Expr, ExpressionMaps}; use extenddb_core::types::{Item, TableKeyInfo}; use extenddb_storage::StreamCapture; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{SortKeyValue, parse_sk, pk_to_text, sk_column, sk_info}; +use extenddb_storage::util::{SortKeyValue, sk_column, sk_info}; + +use super::key_text::{parse_sk, pk_to_text}; use super::index::{enqueue_async_indexes, fetch_write_path_indexes, sync_indexes}; use super::query::check_condition; diff --git a/crates/storage-postgres/src/data/index.rs b/crates/storage-postgres/src/data/index.rs index 551861dbf..07152b2e3 100644 --- a/crates/storage-postgres/src/data/index.rs +++ b/crates/storage-postgres/src/data/index.rs @@ -13,9 +13,11 @@ use extenddb_core::types::{ }; use extenddb_storage::error::StorageError; use extenddb_storage::util::SortKeyValue; -use extenddb_storage::util::{composite_pk_to_text, parse_sk, sk_column, sk_column_n}; +use extenddb_storage::util::{sk_column, sk_column_n}; -use super::{all_sort_key_info, index_table_name}; +use super::key_text::{composite_pk_to_text, parse_sk}; + +use super::{all_sort_key_info, index_table_name, item_to_json}; use crate::gsi_queue::{GsiApplyContext, GsiIndexDef, enqueue_gsi_pending}; /// Map a `sqlx` error to `StorageError`, preserving the SQLSTATE code in the @@ -351,8 +353,7 @@ pub(crate) async fn insert_index_row_multi( let idx_pk_text = composite_pk_to_text(item, index_ks)?; let base_pk_text = composite_pk_to_text(item, base_ks)?; - let item_json = - serde_json::to_value(projected).map_err(|e| StorageError::Internal(e.to_string()))?; + let item_json = item_to_json(projected)?; // Build column list dynamically let mut cols = vec!["pk".to_owned()]; diff --git a/crates/storage-postgres/src/data/key_text.rs b/crates/storage-postgres/src/data/key_text.rs new file mode 100644 index 000000000..a433907af --- /dev/null +++ b/crates/storage-postgres/src/data/key_text.rs @@ -0,0 +1,139 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Key text for the `PostgreSQL` key columns. +//! +//! The shared helpers in `extenddb_storage::util` turn key attribute values into +//! the text the SQL backends bind against `pk`, `sk_s`, `base_pk`, and the other +//! string key columns. `PostgreSQL` `TEXT` cannot hold the byte 0x00, which +//! DynamoDB allows anywhere in a string key, so this module wraps those helpers +//! and passes every string component through the order-preserving escape from +//! `extenddb_storage::util::escape_control`. Numbers and base64 binaries contain +//! neither U+0000 nor U+0001 and are unchanged. The escape preserves byte order +//! and prefixes, so the `COLLATE "C"` comparisons, `BETWEEN`, `begins_with`, +//! and row-comparison cursors that run on these columns return the same rows +//! they would on the raw text. +//! +//! Every site in this crate that produces key column text imports these three +//! functions instead of the shared ones, so the write path, the read path, index +//! rows, cursors, and the propagation queue all agree on the stored form. + +use std::borrow::Cow; + +use extenddb_core::types::{AttributeValue, Item, KeySchemaElement, KeyType, ScalarAttributeType}; +use extenddb_storage::error::StorageError; +use extenddb_storage::util::{self, SortKeyValue, encode_netstring_composite, escape_control}; + +/// Partition key attribute value as escaped column text. +pub(crate) fn pk_to_text(value: &AttributeValue) -> Result, StorageError> { + Ok(match util::pk_to_text(value)? { + Cow::Borrowed(s) => escape_control(s), + Cow::Owned(s) => Cow::Owned(escape_control(&s).into_owned()), + }) +} + +/// Composite partition key as escaped column text. +/// +/// Mirrors `extenddb_storage::util::composite_pk_to_text`: a single HASH +/// attribute is its own text, several are netstring-joined. Each part is +/// escaped before joining, which is also what `query_scan` does when it builds +/// the same text from a key condition, so the two agree byte for byte. +pub(crate) fn composite_pk_to_text( + item: &Item, + key_schema: &[KeySchemaElement], +) -> Result { + let hash_elements: Vec<_> = key_schema + .iter() + .filter(|ks| ks.key_type == KeyType::Hash) + .collect(); + if hash_elements.len() == 1 { + let val = item + .get(&hash_elements[0].attribute_name) + .ok_or_else(|| StorageError::Internal("missing partition key".to_owned()))?; + return Ok(pk_to_text(val)?.into_owned()); + } + let mut parts = Vec::with_capacity(hash_elements.len()); + for ks in &hash_elements { + let val = item.get(&ks.attribute_name).ok_or_else(|| { + StorageError::Internal(format!( + "missing partition key attribute {}", + ks.attribute_name + )) + })?; + parts.push(pk_to_text(val)?.into_owned()); + } + Ok(encode_netstring_composite(&parts)) +} + +/// Sort key attribute value as a typed bind value, string keys escaped. +pub(crate) fn parse_sk( + value: &AttributeValue, + sk_type: ScalarAttributeType, +) -> Result { + Ok(match util::parse_sk(value, sk_type)? { + SortKeyValue::S(s) => SortKeyValue::S(escape_control(&s).into_owned()), + other => other, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn ks(name: &str) -> KeySchemaElement { + KeySchemaElement { + attribute_name: name.to_owned(), + key_type: KeyType::Hash, + } + } + + #[test] + fn string_keys_are_escaped_and_other_types_are_not() { + assert_eq!( + pk_to_text(&AttributeValue::S("a\u{0}b".into())).unwrap(), + "a\u{1}\u{1}b" + ); + assert_eq!( + pk_to_text(&AttributeValue::S("plain".into())).unwrap(), + "plain" + ); + assert_eq!( + pk_to_text(&AttributeValue::N("12.5".into())).unwrap(), + "12.5" + ); + assert_eq!( + pk_to_text(&AttributeValue::B(vec![0, 1, 2])).unwrap(), + util::pk_to_text(&AttributeValue::B(vec![0, 1, 2])).unwrap() + ); + match parse_sk(&AttributeValue::S("\u{1}".into()), ScalarAttributeType::S).unwrap() { + SortKeyValue::S(s) => assert_eq!(s, "\u{1}\u{2}"), + _ => panic!("string sort key"), + } + match parse_sk(&AttributeValue::B(vec![0]), ScalarAttributeType::B).unwrap() { + SortKeyValue::B(b) => assert_eq!(b, vec![0]), + _ => panic!("binary sort key"), + } + } + + #[test] + fn composite_text_escapes_each_part_before_joining() { + let mut item = Item::new(); + item.insert("h1".to_owned(), AttributeValue::S("a\u{0}".into())); + item.insert("h2".to_owned(), AttributeValue::S("b".into())); + let text = composite_pk_to_text(&item, &[ks("h1"), ks("h2")]).unwrap(); + // The escaped part is three bytes, and the netstring length says so. + assert_eq!(text, "3:a\u{1}\u{1},1:b,"); + assert!(!text.contains('\u{0}')); + // Single hash attribute: the escaped value itself. + let text = composite_pk_to_text(&item, &[ks("h1")]).unwrap(); + assert_eq!(text, "a\u{1}\u{1}"); + } + + #[test] + fn stored_text_never_contains_nul() { + for raw in ["\u{0}", "a\u{0}b", "\u{0}\u{1}\u{2}", "x\u{1}\u{1}y"] { + let av = AttributeValue::S(raw.into()); + let text = pk_to_text(&av).unwrap(); + assert!(!text.contains('\u{0}'), "{raw:?} -> {text:?}"); + } + } +} diff --git a/crates/storage-postgres/src/data/mod.rs b/crates/storage-postgres/src/data/mod.rs index d64ceb539..b5ffbc0cc 100755 --- a/crates/storage-postgres/src/data/mod.rs +++ b/crates/storage-postgres/src/data/mod.rs @@ -9,6 +9,7 @@ use extenddb_core::types::{AttributeDefinition, Item, KeySchemaElement, ScalarAttributeType}; use extenddb_storage::error::StorageError; +use extenddb_storage::util::{escape_json_strings, unescape_json_strings}; /// SQL table name for a Virtual `DynamoDB` table. /// @@ -60,8 +61,44 @@ pub(crate) fn all_sort_key_info<'a>( } /// Deserialize an `item_data` JSONB value into an `Item`. +/// +/// Reverses [`item_to_json`]: every string in the stored tree is unescaped +/// before deserialization, so an item written with U+0000 anywhere in it comes +/// back byte-identical. pub(crate) fn json_to_item(v: serde_json::Value) -> Result { - serde_json::from_value(v).map_err(|e| StorageError::Internal(e.to_string())) + serde_json::from_value(unescape_json_strings(v)) + .map_err(|e| StorageError::Internal(e.to_string())) +} + +/// Serialize an `Item` for an `item_data` JSONB column. +/// +/// PostgreSQL `jsonb` rejects the `\u0000` escape, so every string in the tree +/// (attribute names, map keys, string values) goes through the order-preserving +/// escape from `extenddb_storage::util` before it reaches the column. Items +/// without U+0000 or U+0001 serialize to exactly what they did before. +pub(crate) fn item_to_json(item: &Item) -> Result { + serde_json::to_value(item) + .map(escape_json_strings) + .map_err(|e| StorageError::Internal(e.to_string())) +} + +/// Serialize any value for a JSONB column that may carry item strings (stream +/// records, queued index updates, index contexts). Same escape as +/// [`item_to_json`]; the inverse is [`stored_json_to`]. +pub(crate) fn to_stored_json( + value: &T, +) -> Result { + serde_json::to_value(value) + .map(escape_json_strings) + .map_err(|e| StorageError::Internal(e.to_string())) +} + +/// Deserialize a JSONB value written by [`to_stored_json`]. +pub(crate) fn stored_json_to( + v: serde_json::Value, +) -> Result { + serde_json::from_value(unescape_json_strings(v)) + .map_err(|e| StorageError::Internal(e.to_string())) } /// Bind a `SortKeyValue` to a positional parameter in a sqlx query and execute it. @@ -133,6 +170,7 @@ mod data_engine; mod ddl; mod delete_item; pub(crate) mod index; +pub(crate) mod key_text; mod put_item; mod query; mod query_scan; diff --git a/crates/storage-postgres/src/data/put_item.rs b/crates/storage-postgres/src/data/put_item.rs index e307df660..6ce8ec2f2 100755 --- a/crates/storage-postgres/src/data/put_item.rs +++ b/crates/storage-postgres/src/data/put_item.rs @@ -7,12 +7,14 @@ use extenddb_core::expression::{Expr, ExpressionMaps}; use extenddb_core::types::{Item, TableKeyInfo}; use extenddb_storage::StreamCapture; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{composite_pk_to_text, parse_sk, pk_to_text, sk_column, sk_info}; +use extenddb_storage::util::{sk_column, sk_info}; + +use super::key_text::{composite_pk_to_text, parse_sk, pk_to_text}; use super::index::{enqueue_async_indexes, fetch_write_path_indexes, sync_indexes}; use super::query::check_condition; use super::tx_helpers::write_stream_record_in_tx; -use super::{data_table_name, json_to_item}; +use super::{data_table_name, item_to_json, json_to_item}; use crate::PostgresEngine; impl PostgresEngine { @@ -30,8 +32,7 @@ impl PostgresEngine { let pk_text = composite_pk_to_text(&item, &key_info.key_schema)?; - let item_json = - serde_json::to_value(&item).map_err(|e| StorageError::Internal(e.to_string()))?; + let item_json = item_to_json(&item)?; // Both index families in one catalog visit (D-4: sync + async split for the // secondary indexes). diff --git a/crates/storage-postgres/src/data/query.rs b/crates/storage-postgres/src/data/query.rs index 6a84dc1b7..ca3c8ddcc 100755 --- a/crates/storage-postgres/src/data/query.rs +++ b/crates/storage-postgres/src/data/query.rs @@ -12,7 +12,9 @@ use extenddb_core::types::{ }; use extenddb_storage::error::StorageError; use extenddb_storage::util::SortKeyValue; -use extenddb_storage::util::{parse_sk, pk_to_text, sk_info}; +use extenddb_storage::util::sk_info; + +use super::key_text::{parse_sk, pk_to_text}; /// Extra pagination bind values for index queries. /// diff --git a/crates/storage-postgres/src/data/query_scan.rs b/crates/storage-postgres/src/data/query_scan.rs index c0f0c8b05..ef4525f9c 100644 --- a/crates/storage-postgres/src/data/query_scan.rs +++ b/crates/storage-postgres/src/data/query_scan.rs @@ -6,9 +6,9 @@ use extenddb_core::expression::{ExpressionMaps, KeyCondition}; use extenddb_core::types::{Item, ScalarAttributeType, TableKeyInfo}; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{ - encode_netstring_composite, parse_sk, pk_to_text, sk_column, sk_column_n, sk_info, -}; +use extenddb_storage::util::{encode_netstring_composite, sk_column, sk_column_n, sk_info}; + +use super::key_text::{parse_sk, pk_to_text}; use super::query::{ PaginationBinds, build_key, build_sk_sql, execute_query_sql, execute_scan_sql, diff --git a/crates/storage-postgres/src/data/tx_helpers.rs b/crates/storage-postgres/src/data/tx_helpers.rs index ed30a1bf3..aa1454555 100755 --- a/crates/storage-postgres/src/data/tx_helpers.rs +++ b/crates/storage-postgres/src/data/tx_helpers.rs @@ -12,9 +12,11 @@ use extenddb_core::types::{ }; use extenddb_storage::StreamCapture; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{SortKeyValue, parse_sk, pk_to_text, sk_column, sk_info}; +use extenddb_storage::util::{SortKeyValue, sk_column, sk_info}; -use super::{data_table_name, json_to_item}; +use super::key_text::{parse_sk, pk_to_text}; + +use super::{data_table_name, item_to_json, json_to_item, to_stored_json}; /// Fetch a single item within an existing transaction. pub(super) async fn fetch_item_in_tx( @@ -105,8 +107,7 @@ pub(super) async fn upsert_item_in_tx( .get(pk_name) .ok_or_else(|| StorageError::Internal("missing partition key".to_owned()))?; let pk_text = pk_to_text(pk_value)?; - let item_json = - serde_json::to_value(item).map_err(|e| StorageError::Internal(e.to_string()))?; + let item_json = item_to_json(item)?; if let Some((sk_name, sk_type)) = sk_info(&key_info.key_schema, &key_info.attribute_definitions) { @@ -157,8 +158,7 @@ pub(super) async fn insert_item_if_absent_in_tx( .get(pk_name) .ok_or_else(|| StorageError::Internal("missing partition key".to_owned()))?; let pk_text = pk_to_text(pk_value)?; - let item_json = - serde_json::to_value(item).map_err(|e| StorageError::Internal(e.to_string()))?; + let item_json = item_to_json(item)?; let rows_affected = if let Some((sk_name, sk_type)) = sk_info(&key_info.key_schema, &key_info.attribute_definitions) @@ -366,8 +366,7 @@ pub(super) async fn write_stream_record_in_tx( user_identity: capture.user_identity.clone(), }; - let record_json = - serde_json::to_value(&record).map_err(|e| StorageError::Internal(e.to_string()))?; + let record_json = to_stored_json(&record)?; sqlx::query( "INSERT INTO stream_records (sequence_number, shard_id, table_id, event_name, record_data) \ diff --git a/crates/storage-postgres/src/data/update_item.rs b/crates/storage-postgres/src/data/update_item.rs index 72ce6edb9..7a731309d 100755 --- a/crates/storage-postgres/src/data/update_item.rs +++ b/crates/storage-postgres/src/data/update_item.rs @@ -8,9 +8,12 @@ use extenddb_core::types::{Item, TableKeyInfo}; use extenddb_core::validation; use extenddb_storage::StreamCapture; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{parse_sk, pk_to_text, sk_column, sk_info}; +use extenddb_storage::util::{sk_column, sk_info}; + +use super::key_text::{parse_sk, pk_to_text}; use super::index::{enqueue_async_indexes, fetch_write_path_indexes, sync_indexes}; +use super::item_to_json; use super::query::check_condition; use super::tx_helpers::write_stream_record_in_tx; use super::{data_table_name, json_to_item}; @@ -192,8 +195,7 @@ impl PostgresEngine { let new_item = if return_new { Some(item.clone()) } else { None }; - let item_json = - serde_json::to_value(&item).map_err(|e| StorageError::Internal(e.to_string()))?; + let item_json = item_to_json(&item)?; if old_json.is_some() { // Row existed and is locked by the read above, so update in place. diff --git a/crates/storage-postgres/src/data/vector_index.rs b/crates/storage-postgres/src/data/vector_index.rs index 0bc0a18b5..500225ed8 100644 --- a/crates/storage-postgres/src/data/vector_index.rs +++ b/crates/storage-postgres/src/data/vector_index.rs @@ -32,15 +32,15 @@ use extenddb_core::types::{AttributeDefinition, Item, KeySchemaElement, ScalarAttributeType}; use extenddb_core::validation::vector_item::{vector_components, vector_norm}; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{ - SortKeyValue, composite_pk_to_text, parse_sk, sk_column, sk_column_n, -}; +use extenddb_storage::util::{SortKeyValue, sk_column, sk_column_n}; + +use super::key_text::{composite_pk_to_text, parse_sk}; use extenddb_storage::vector_lifecycle::{ VectorApplyContext, VectorIndexMeta, item_is_indexable, item_partition, projected_payload, }; use pgvector::Vector; -use super::{all_sort_key_info, vector_table_name}; +use super::{all_sort_key_info, item_to_json, vector_table_name}; /// Read the vector index metadata for a table from the catalog. /// @@ -403,8 +403,7 @@ impl<'a> VectorInsertPlan<'a> { // vector attribute are the shared payload rules, so a live-written row and a // backfilled one cannot differ in shape. let projected = projected_payload(item, base_key_schema, meta); - let item_json = - serde_json::to_value(&projected).map_err(|e| StorageError::Internal(e.to_string()))?; + let item_json = item_to_json(&projected)?; // A plain INSERT, deliberately, where the GSI sibling upserts. Two callers reach // it and only one deletes first: the live write path through diff --git a/crates/storage-postgres/src/gsi_queue.rs b/crates/storage-postgres/src/gsi_queue.rs index 39e60d443..ab567dbb2 100644 --- a/crates/storage-postgres/src/gsi_queue.rs +++ b/crates/storage-postgres/src/gsi_queue.rs @@ -23,9 +23,9 @@ use std::sync::Arc; +use crate::data::key_text::composite_pk_to_text; use extenddb_core::types::{AttributeDefinition, Item, KeySchemaElement, Projection}; use extenddb_storage::error::StorageError; -use extenddb_storage::util::composite_pk_to_text; use serde::{Deserialize, Serialize}; use sqlx::PgPool; use tokio::sync::Notify; @@ -223,16 +223,9 @@ pub(crate) async fn enqueue_gsi_pending( delay_ms: u64, context: &PendingApplyContext, ) -> Result<(), StorageError> { - let old_json = old_item - .map(serde_json::to_value) - .transpose() - .map_err(|e| StorageError::Internal(e.to_string()))?; - let new_json = new_item - .map(serde_json::to_value) - .transpose() - .map_err(|e| StorageError::Internal(e.to_string()))?; - let context_json = - serde_json::to_value(context).map_err(|e| StorageError::Internal(e.to_string()))?; + let old_json = old_item.map(crate::data::item_to_json).transpose()?; + let new_json = new_item.map(crate::data::item_to_json).transpose()?; + let context_json = crate::data::to_stored_json(context)?; // Route all updates for a given base item to one worker (per-key FIFO). // The base key is immutable over an item's lifetime; `new_item` carries it @@ -426,16 +419,9 @@ async fn apply_claimed_row( new_json: Option, ctx_json: serde_json::Value, ) -> Result<(), StorageError> { - let old_item: Option = old_json - .map(serde_json::from_value) - .transpose() - .map_err(|e| StorageError::Internal(e.to_string()))?; - let new_item: Option = new_json - .map(serde_json::from_value) - .transpose() - .map_err(|e| StorageError::Internal(e.to_string()))?; - let context: PendingApplyContext = - serde_json::from_value(ctx_json).map_err(|e| StorageError::Internal(e.to_string()))?; + let old_item: Option = old_json.map(crate::data::json_to_item).transpose()?; + let new_item: Option = new_json.map(crate::data::json_to_item).transpose()?; + let context: PendingApplyContext = crate::data::stored_json_to(ctx_json)?; // One index per row. Guard the apply with a savepoint so a dropped-index // race (42P01) can be recovered and the row still consumed; aborting the diff --git a/crates/storage-postgres/src/lib.rs b/crates/storage-postgres/src/lib.rs index 0d7a51c0d..3dfeb03b3 100755 --- a/crates/storage-postgres/src/lib.rs +++ b/crates/storage-postgres/src/lib.rs @@ -406,6 +406,38 @@ impl PostgresEngine { Ok(()) } + /// Refuse to serve a data database whose stored strings have not been + /// brought to this build's encoding. + /// + /// The escape in `extenddb_storage::util::escape_control` stores U+0001 as a + /// two-character sequence, and the migration that re-encodes rows written by + /// older builds (`004_escape_control_chars`) rewrites every row that contains + /// U+0001. It cannot tell a row this build wrote from a legacy one, so this + /// build must not write anything before the migration has run. `extenddb + /// init` records the migration on a fresh deployment; an upgrade runs + /// `extenddb migrate --yes` after stopping every older server and before + /// starting this one. + /// + /// # Errors + /// + /// Returns `StorageError::Internal` naming the migration and the command to + /// run when the data database's `schema_history` does not record it. + pub async fn check_data_migrations_applied(&self) -> Result<(), StorageError> { + let missing = migrations::unapplied_required_data_migrations(&self.data_pool) + .await + .map_err(|e| StorageError::Connection(format!("{e:?}")))?; + if missing.is_empty() { + return Ok(()); + } + Err(StorageError::Internal(format!( + "Data migration not applied: {}. This build stores strings in a different \ + encoding from earlier builds, so it must not serve until the existing rows \ + have been re-encoded. Stop every older ExtendDB server that uses this data \ + database, run 'extenddb migrate --yes', then start this server.", + missing.join(", ") + ))) + } + /// Query the data database name from the catalog for the startup banner (REQ-LOG-001). /// /// Returns `"(not configured)"` if no data database has been registered. @@ -656,6 +688,18 @@ fn server_components_factory( _ => BackendError::InitializationFailed(e.to_string()), })?; + // This build stores strings through the control-character escape, and the + // re-encoding migration cannot tell a row this build wrote from a row an + // older build wrote (both can hold U+0001). Serving before the migration + // has run would let the migration later rewrite this build's own rows, so + // refuse to start until it is recorded. `extenddb init` records it on a + // fresh deployment; an upgraded deployment runs `extenddb migrate --yes` + // with every older server stopped. + engine + .check_data_migrations_applied() + .await + .map_err(|e| BackendError::InitializationFailed(e.to_string()))?; + // Rebuild any vector index a crash left CREATING, before serving: an index // in that state is not searchable, and nothing else will repair it. match crate::data::vector_index::reconcile_incomplete_vector_indexes(&engine).await { diff --git a/crates/storage-postgres/src/metadata_engine.rs b/crates/storage-postgres/src/metadata_engine.rs index 322a9015f..bd553b3e9 100755 --- a/crates/storage-postgres/src/metadata_engine.rs +++ b/crates/storage-postgres/src/metadata_engine.rs @@ -273,7 +273,9 @@ impl MetadataEngine for PostgresEngine { ) -> BoxFuture<'_, Result<(), StorageError>> { let account_id = account_id.to_string(); let table_name = table_name.to_string(); - let ttl_attribute = ttl_attribute.to_string(); + // Stored attribute names carry the same escape as the rest of item_data, + // so the name SQL addresses must be the stored form. + let ttl_attribute = extenddb_storage::util::escape_control(ttl_attribute).into_owned(); Box::pin(async move { Self::validate_account_id(&account_id)?; let (table_id,): (String,) = sqlx::query_as( @@ -366,7 +368,9 @@ impl MetadataEngine for PostgresEngine { ) -> BoxFuture<'_, Result, StorageError>> { let account_id = account_id.to_string(); let table_name = table_name.to_string(); - let ttl_attribute = ttl_attribute.to_string(); + // Stored attribute names carry the same escape as the rest of item_data, + // so the name SQL addresses must be the stored form. + let ttl_attribute = extenddb_storage::util::escape_control(ttl_attribute).into_owned(); Box::pin(async move { Self::validate_account_id(&account_id)?; let (table_id,): (String,) = sqlx::query_as( diff --git a/crates/storage-postgres/src/migrations.rs b/crates/storage-postgres/src/migrations.rs index 45af16167..eaa5ffb36 100755 --- a/crates/storage-postgres/src/migrations.rs +++ b/crates/storage-postgres/src/migrations.rs @@ -129,7 +129,91 @@ pub(crate) async fn run_data_migrations(pool: &PgPool) -> OpResult<()> { /// outside a transaction (they use `CREATE INDEX CONCURRENTLY`), so they cannot /// be expressed as SQL in [`DATA_MIGRATIONS`]. Applied by `extenddb migrate` /// after the SQL migrations, so the operator controls when the change happens. -pub(crate) const DATA_CODE_MIGRATIONS: &[&str] = &["003_gsi_base_key_index"]; +pub(crate) const DATA_CODE_MIGRATIONS: &[&str] = + &["003_gsi_base_key_index", "004_escape_control_chars"]; + +/// Code migrations a server must see recorded before it serves. `003` only adds +/// an index and a server runs correctly without it; `004` changes the stored +/// form of strings, and a server that wrote through the new escape before the +/// migration ran would have its own rows rewritten by it (see +/// [`escape_legacy_control_chars`]). `PostgresEngine::check_data_migrations_applied` +/// refuses to start on a data database missing any of these. +pub(crate) const REQUIRED_DATA_CODE_MIGRATIONS: &[&str] = &["004_escape_control_chars"]; + +/// `application_name` the init and migrate commands set on their connections, +/// so [`refuse_if_other_clients_connected`] can leave them out of its count. +pub(crate) const MIGRATE_APPLICATION_NAME: &str = "extenddb-migrate"; + +/// Environment variable that skips [`refuse_if_other_clients_connected`], for a +/// data database whose only other connections are known to be idle (a +/// connection pooler holding server-side connections open, for example). +pub(crate) const IGNORE_CONNECTIONS_ENV: &str = "EXTENDDB_MIGRATE_IGNORE_CONNECTIONS"; + +/// Refuse to run a data-rewriting migration while anything else is connected +/// to the data database. +/// +/// `004_escape_control_chars` rewrites rows in place and cannot tell a row a +/// running server writes from a legacy row: an older server still writing +/// leaves unescaped rows behind, and a server of this release writing before +/// the migration finishes gets its rows double-escaped. The upgrade manual +/// requires every server to be stopped first; this check catches the case +/// where one was not. Connections carrying [`MIGRATE_APPLICATION_NAME`] are +/// this process's own and are ignored; every other client backend on the +/// database counts. Managed environments where a pooler keeps idle +/// connections open can set [`IGNORE_CONNECTIONS_ENV`] to skip the check. +pub(crate) async fn refuse_if_other_clients_connected(pool: &PgPool) -> OpResult<()> { + let others: Vec<(String, Option, Option)> = sqlx::query_as( + "SELECT usename::text, application_name, client_addr::text \ + FROM pg_stat_activity \ + WHERE datname = current_database() \ + AND backend_type = 'client backend' \ + AND pid <> pg_backend_pid() \ + AND coalesce(application_name, '') <> $1", + ) + .bind(MIGRATE_APPLICATION_NAME) + .fetch_all(pool) + .await + .map_err(|e| OpError::Internal(format!("Check for other connections: {e}")))?; + if others.is_empty() { + return Ok(()); + } + let mut described: Vec = others + .iter() + .map(|(user, app, addr)| { + format!( + "{user} from {} ({})", + addr.as_deref().unwrap_or("local socket"), + match app.as_deref() { + Some("") | None => "no application name", + Some(a) => a, + } + ) + }) + .collect(); + described.sort(); + described.dedup(); + Err(OpError::Internal(format!( + "{} other connection(s) hold the data database: {}. This migration rewrites \ + stored rows and must run with every ExtendDB server stopped and every other \ + session closed. Stop them and run 'extenddb migrate --yes' again. If the only \ + other connections are idle ones held by a connection pooler, set \ + {IGNORE_CONNECTIONS_ENV}=1 to skip this check.", + others.len(), + described.join("; ") + ))) +} + +/// Names in [`REQUIRED_DATA_CODE_MIGRATIONS`] not recorded in this data +/// database's `schema_history`. +pub(crate) async fn unapplied_required_data_migrations(pool: &PgPool) -> OpResult> { + let mut missing = Vec::new(); + for name in REQUIRED_DATA_CODE_MIGRATIONS { + if !is_migration_applied(pool, name).await? { + missing.push((*name).to_owned()); + } + } + Ok(missing) +} /// Run programmatic data migrations, skipping already-applied ones. /// @@ -151,6 +235,15 @@ pub(crate) async fn run_data_code_migrations( "003_gsi_base_key_index" => { ensure_gsi_base_key_indexes(catalog_pool, data_pool).await?; } + "004_escape_control_chars" => { + // Rows written by a server that is still running would be missed + // (an older release) or double-escaped (this release); see + // `refuse_if_other_clients_connected`. + if std::env::var_os(IGNORE_CONNECTIONS_ENV).is_none() { + refuse_if_other_clients_connected(data_pool).await?; + } + escape_legacy_control_chars(catalog_pool, data_pool).await?; + } other => { return Err(OpError::Internal(format!( "Unknown data code migration: {other}" @@ -219,6 +312,355 @@ async fn ensure_gsi_base_key_indexes(catalog_pool: &PgPool, data_pool: &PgPool) Ok(()) } +/// Progress ledger for multi-table data code migrations: one row per +/// (migration, table) pair, committed in the same transaction as that table's +/// rewrite, so a re-run after a crash knows exactly which tables are done. +const CODE_MIGRATION_PROGRESS_TABLE: &str = "data_code_migration_progress"; + +/// Create the progress ledger if this database does not have one yet. +async fn ensure_progress_table(pool: &PgPool) -> OpResult<()> { + sqlx::query(&format!( + "CREATE TABLE IF NOT EXISTS {CODE_MIGRATION_PROGRESS_TABLE} (\ + migration_name TEXT NOT NULL, \ + table_name TEXT NOT NULL, \ + completed_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), \ + PRIMARY KEY (migration_name, table_name)\ + )" + )) + .execute(pool) + .await + .map_err(|e| OpError::Internal(format!("Create migration progress table: {e}")))?; + Ok(()) +} + +/// Whether the progress ledger says `table` is already done for `migration`. +async fn is_table_marked(pool: &PgPool, migration: &str, table: &str) -> OpResult { + let marked: bool = sqlx::query_scalar(&format!( + "SELECT EXISTS(SELECT 1 FROM {CODE_MIGRATION_PROGRESS_TABLE} \ + WHERE migration_name = $1 AND table_name = $2)" + )) + .bind(migration) + .bind(table) + .fetch_one(pool) + .await + .map_err(|e| OpError::Internal(format!("Check migration progress: {e}")))?; + Ok(marked) +} + +/// One physical table's rewrite work for `004_escape_control_chars`: the text +/// key columns that get the single-statement SQL replace, and the jsonb +/// columns whose rows are re-encoded one by one. +struct EscapeWork { + /// Bare (unquoted) table name, e.g. `_ddb_` or `gsi_pending`. + table: String, + text_columns: Vec, + jsonb_columns: Vec, +} + +/// Rewrite work for a dynamically named `_ddb_*` table, read from +/// `information_schema.columns` so that exactly the sort key columns this +/// table has (`sk_s`, `sk_s`, `base_sk_s`, ...) are covered. Returns +/// `None` when the physical table does not exist (a catalog row mid-create or +/// mid-delete): such a table has no legacy rows to re-encode. +async fn dynamic_table_work(pool: &PgPool, table: String) -> OpResult> { + let text_columns: Vec = sqlx::query_scalar( + "SELECT column_name::text FROM information_schema.columns \ + WHERE table_schema = 'public' AND table_name = $1 AND data_type = 'text' \ + AND (column_name IN ('pk', 'base_pk') OR column_name::text ~ '^(base_)?sk[0-9]*_s$') \ + ORDER BY ordinal_position", + ) + .bind(&table) + .fetch_all(pool) + .await + .map_err(|e| OpError::Internal(format!("List text key columns of {table}: {e}")))?; + + let jsonb_columns: Vec = sqlx::query_scalar( + "SELECT column_name::text FROM information_schema.columns \ + WHERE table_schema = 'public' AND table_name = $1 AND data_type = 'jsonb' \ + AND column_name = 'item_data'", + ) + .bind(&table) + .fetch_all(pool) + .await + .map_err(|e| OpError::Internal(format!("List jsonb columns of {table}: {e}")))?; + + if text_columns.is_empty() && jsonb_columns.is_empty() { + return Ok(None); + } + Ok(Some(EscapeWork { + table, + text_columns, + jsonb_columns, + })) +} + +/// Rewrite work for a fixed-name table, with the candidate columns filtered +/// through `information_schema.columns` so a column absent on this deployment +/// is skipped. Returns `None` when the table itself is absent. +async fn fixed_table_work( + pool: &PgPool, + table: &str, + text_candidates: &[&str], + jsonb_candidates: &[&str], +) -> OpResult> { + if !table_exists(pool, table).await? { + return Ok(None); + } + let text_candidates: Vec = text_candidates.iter().map(|c| (*c).to_owned()).collect(); + let jsonb_candidates: Vec = jsonb_candidates.iter().map(|c| (*c).to_owned()).collect(); + let text_columns: Vec = sqlx::query_scalar( + "SELECT column_name::text FROM information_schema.columns \ + WHERE table_schema = 'public' AND table_name = $1 AND data_type = 'text' \ + AND column_name::text = ANY($2) \ + ORDER BY ordinal_position", + ) + .bind(table) + .bind(&text_candidates) + .fetch_all(pool) + .await + .map_err(|e| OpError::Internal(format!("List text columns of {table}: {e}")))?; + let jsonb_columns: Vec = sqlx::query_scalar( + "SELECT column_name::text FROM information_schema.columns \ + WHERE table_schema = 'public' AND table_name = $1 AND data_type = 'jsonb' \ + AND column_name::text = ANY($2) \ + ORDER BY ordinal_position", + ) + .bind(table) + .bind(&jsonb_candidates) + .fetch_all(pool) + .await + .map_err(|e| OpError::Internal(format!("List jsonb columns of {table}: {e}")))?; + Ok(Some(EscapeWork { + table: table.to_owned(), + text_columns, + jsonb_columns, + })) +} + +/// The re-encoded form of a stored jsonb value, or `None` when re-encoding +/// does not change it. `escape_json_strings` only changes strings and object +/// keys containing U+0000 or U+0001, so a plain document and a document whose +/// strings merely contain the literal six-character text backslash-u-0001 both +/// come back equal and are skipped. +fn reencoded_json(value: &serde_json::Value) -> Option { + let escaped = extenddb_storage::util::escape_json_strings(value.clone()); + (escaped != *value).then_some(escaped) +} + +/// Re-encode the rows of one jsonb column whose stored document contains +/// U+0001. Runs inside the caller's per-table transaction. +/// +/// The candidate ctids are snapshotted once, before any row of this column is +/// touched, so no row is ever visited twice; the updates relocate rows (a ctid +/// names a physical row version), but always after that row's value was read. +/// A parameterized `DECLARE CURSOR` cannot be issued through the extended +/// protocol, so the locked-snapshot form of a streamed fetch is used instead, +/// and the values are then fetched in batches of 500 by ctid so a large table +/// never loads into memory. The matching set is tiny in practice; the table +/// scan, not the writes, is the cost. +async fn escape_jsonb_rows( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + quoted_table: &str, + col: &str, +) -> OpResult { + // '\u0001' (six characters) is how jsonb::text renders a raw U+0001 inside + // a string or key. It also matches a stored string whose content is a + // literal backslash followed by u0001 (rendered '\\u0001'), which is + // harmless: that row re-encodes to itself and is skipped below. jsonb text + // rendering is why SQL-level replace() is not used here: it would corrupt + // exactly that literal-backslash case. The backslash is spelled chr(92) so + // the predicate reads the same under either value of + // standard_conforming_strings; a plain '\u0001' literal would become the + // single byte 0x01 with the setting off and match nothing. + let select_ctids = format!( + "SELECT ctid::text FROM {quoted_table} \ + WHERE position((chr(92) || 'u0001') IN \"{col}\"::text) > 0 FOR UPDATE" + ); + let ctids: Vec = sqlx::query_scalar(&select_ctids) + .fetch_all(&mut **tx) + .await + .map_err(|e| OpError::Internal(format!("Find rows to re-encode in {quoted_table}: {e}")))?; + + let fetch_batch = + format!("SELECT ctid::text, \"{col}\" FROM {quoted_table} WHERE ctid = ANY($1::tid[])"); + let update_row = format!("UPDATE {quoted_table} SET \"{col}\" = $1 WHERE ctid = $2::tid"); + let mut changed: u64 = 0; + for batch in ctids.chunks(500) { + let rows: Vec<(String, serde_json::Value)> = sqlx::query_as(&fetch_batch) + .bind(batch.to_vec()) + .fetch_all(&mut **tx) + .await + .map_err(|e| OpError::Internal(format!("Fetch rows from {quoted_table}: {e}")))?; + for (ctid, value) in rows { + if let Some(escaped) = reencoded_json(&value) { + sqlx::query(&update_row) + .bind(escaped) + .bind(&ctid) + .execute(&mut **tx) + .await + .map_err(|e| { + OpError::Internal(format!("Re-encode a row of {quoted_table}: {e}")) + })?; + changed += 1; + } + } + } + Ok(changed) +} + +/// Rewrite one table in one transaction and mark it done in the progress +/// ledger before committing, so the rewrite and the marker are atomic. +async fn escape_rows_in_table(pool: &PgPool, work: &EscapeWork) -> OpResult { + let mut tx = pool + .begin() + .await + .map_err(|e| OpError::Internal(format!("Begin transaction for {}: {e}", work.table)))?; + let quoted = format!("\"{}\"", work.table); + let mut rewritten: u64 = 0; + + // Key text columns: TEXT has no escaping layer, so a single SQL statement + // per column is exact. NULL columns never match (position() is NULL). + for col in &work.text_columns { + let sql = format!( + "UPDATE {quoted} SET \"{col}\" = replace(\"{col}\", chr(1), chr(1) || chr(2)) \ + WHERE position(chr(1) IN \"{col}\") > 0" + ); + let done = sqlx::query(&sql) + .execute(&mut *tx) + .await + .map_err(|e| OpError::Internal(format!("Re-encode {}.{col}: {e}", work.table)))?; + rewritten += done.rows_affected(); + } + + for col in &work.jsonb_columns { + rewritten += escape_jsonb_rows(&mut tx, "ed, col).await?; + } + + sqlx::query(&format!( + "INSERT INTO {CODE_MIGRATION_PROGRESS_TABLE} (migration_name, table_name) \ + VALUES ($1, $2) ON CONFLICT DO NOTHING" + )) + .bind("004_escape_control_chars") + .bind(&work.table) + .execute(&mut *tx) + .await + .map_err(|e| OpError::Internal(format!("Mark {} as re-encoded: {e}", work.table)))?; + + tx.commit() + .await + .map_err(|e| OpError::Internal(format!("Commit re-encode of {}: {e}", work.table)))?; + Ok(rewritten) +} + +/// Re-encode legacy rows that contain a raw U+0001 (`004_escape_control_chars`). +/// +/// Strings are stored through the order-preserving escape in +/// `extenddb_storage::util::escape_control` (U+0000 becomes U+0001 U+0001, +/// U+0001 becomes U+0001 U+0002). Rows written before the escape existed hold +/// their strings raw. The decoder keeps a stray U+0001 literal, so the only +/// legacy rows it can misread are those where a raw U+0001 is immediately +/// followed by U+0001 or U+0002; rewriting every row containing U+0001 removes +/// that case. Legacy rows never contain U+0000 (PostgreSQL refused it), so for +/// legacy data applying the escape to the raw stored value is exactly the +/// re-encode, and rows without U+0001 encode to themselves and are left alone. +/// +/// Crash safety. This migration is NOT idempotent at the row level: escaped +/// data also contains U+0001 (the first character of every escape pair), and a +/// legacy U+0001 U+0002 is indistinguishable from an escaped U+0001, so +/// applying the escape to an already-escaped row would double-escape it. A +/// partially rewritten table must therefore never be rescanned. Each table is +/// processed in one transaction: its matching rows are rewritten, a marker row +/// (migration_name, table_name) is inserted into the small progress table +/// created at the start of this migration, and the transaction commits. A +/// re-run after a crash skips tables that carry a marker and resumes with the +/// rest; when every table is marked, the caller records the migration in +/// `schema_history` exactly as `003` does, and the usual +/// `is_migration_applied` guard makes later runs no-ops. Marker rows are kept +/// forever: deleting them before the `schema_history` row is durable would +/// reopen the rescan window. +/// +/// Rows written by THIS build are indistinguishable from legacy rows for the +/// same reason (an escaped U+0001 is stored as U+0001 U+0002), so the server +/// refuses to start on a data database that does not record this migration +/// (`PostgresEngine::check_data_migrations_applied`). The migration therefore +/// only ever sees rows written by older builds, provided every older server +/// is stopped before it runs; `extenddb init` records it on a fresh deployment. `backup_items` lives in the catalog database +/// while every other table lives in the data database, and a transaction +/// cannot span two databases, so each database gets its own progress table and +/// each marker commits with the table it marks. +async fn escape_legacy_control_chars(catalog_pool: &PgPool, data_pool: &PgPool) -> OpResult<()> { + const NAME: &str = "004_escape_control_chars"; + ensure_progress_table(data_pool).await?; + ensure_progress_table(catalog_pool).await?; + + let mut work: Vec<(&PgPool, EscapeWork)> = Vec::new(); + + // Every data table, index table, and vector index table, enumerated from + // the catalog. Their text key columns are read from information_schema so + // single-sort-key, multi-sort-key, and base_* layouts are all covered. + let table_ids: Vec = + sqlx::query_scalar("SELECT table_id FROM tables ORDER BY table_id") + .fetch_all(catalog_pool) + .await + .map_err(|e| OpError::Internal(format!("Enumerate tables: {e}")))?; + for id in table_ids { + if let Some(w) = dynamic_table_work(data_pool, format!("_ddb_{id}")).await? { + work.push((data_pool, w)); + } + } + let index_ids: Vec = + sqlx::query_scalar("SELECT index_id FROM indexes ORDER BY index_id") + .fetch_all(catalog_pool) + .await + .map_err(|e| OpError::Internal(format!("Enumerate indexes: {e}")))?; + for id in index_ids { + if let Some(w) = dynamic_table_work(data_pool, format!("_ddb_{id}")).await? { + work.push((data_pool, w)); + } + } + let vector_ids: Vec = + sqlx::query_scalar("SELECT index_id FROM vector_indexes ORDER BY index_id") + .fetch_all(catalog_pool) + .await + .map_err(|e| OpError::Internal(format!("Enumerate vector indexes: {e}")))?; + for id in vector_ids { + if let Some(w) = dynamic_table_work(data_pool, format!("_ddb_vec_{id}")).await? { + work.push((data_pool, w)); + } + } + + // Fixed tables holding item documents. gsi_pending and stream_records live + // in the data database; backup_items lives in the catalog database. + if let Some(w) = fixed_table_work( + data_pool, + "gsi_pending", + &[], + &["old_item", "new_item", "index_context"], + ) + .await? + { + work.push((data_pool, w)); + } + if let Some(w) = fixed_table_work(data_pool, "stream_records", &[], &["record_data"]).await? { + work.push((data_pool, w)); + } + if let Some(w) = + fixed_table_work(catalog_pool, "backup_items", &["pk", "sk"], &["item_data"]).await? + { + work.push((catalog_pool, w)); + } + + for (pool, w) in &work { + if is_table_marked(pool, NAME, &w.table).await? { + println!(" {}: already re-encoded, skipping.", w.table); + continue; + } + let rewritten = escape_rows_in_table(pool, w).await?; + println!(" {}: {rewritten} row update(s)", w.table); + } + Ok(()) +} + /// Check if a table exists in the public schema. pub(crate) async fn table_exists(pool: &PgPool, name: &str) -> OpResult { let exists: bool = sqlx::query_scalar( @@ -299,6 +741,30 @@ mod tests { use super::CATALOG_MIGRATIONS; use crate::CATALOG_VERSION; + /// `reencoded_json` is what decides whether a stored jsonb row is + /// rewritten by `004_escape_control_chars`, so it must change exactly the + /// rows holding a raw control character and nothing else. + #[test] + fn reencoded_json_changes_exactly_the_rows_holding_raw_control_chars() { + use serde_json::json; + + // Raw U+0001 in a string value and in an object key: rewritten, with + // each U+0001 becoming U+0001 U+0002 wherever it appears. + let raw = json!({"k\u{1}": {"S": "v\u{1}"}, "plain": {"N": "1"}}); + let expected = json!({"k\u{1}\u{2}": {"S": "v\u{1}\u{2}"}, "plain": {"N": "1"}}); + assert_eq!(super::reencoded_json(&raw), Some(expected)); + + // A string whose content is the literal six-character text + // backslash-u-0001: the SQL candidate filter over-matches it, and this + // check is what keeps it byte-identical. + let literal = json!({"note": {"S": "x\\u0001y"}}); + assert_eq!(super::reencoded_json(&literal), None); + + // A plain document: untouched. + let plain = json!({"pk": {"S": "p"}, "m": {"M": {"k": {"S": "v"}}}}); + assert_eq!(super::reencoded_json(&plain), None); + } + /// The catalog version and the migration list must move together. /// /// A migration that creates its schema without moving the version leaves a @@ -361,3 +827,572 @@ mod tests { } } } + +#[cfg(test)] +mod live_reencode { + //! Live checks for `004_escape_control_chars` against a scratch database. + //! + //! Follows the convention of `tests/key_collation.rs`: each test needs + //! `EXTENDDB_TEST_PG_CONNECTION_STRING` (a base URL with no database name, + //! for example `postgresql://postgres:postgres@127.0.0.1:5432`), builds a + //! throwaway database, and drops it when it passes. Without the variable + //! every test here reports a skip and passes. One database serves as both + //! catalog and data, as the storage-level tests do. + + use serde_json::json; + use sqlx::PgPool; + use sqlx::postgres::PgPoolOptions; + + use super::{ + CATALOG_MIGRATIONS, DATA_MIGRATIONS, IGNORE_CONNECTIONS_ENV, MIGRATE_APPLICATION_NAME, + pending_data_migrations, refuse_if_other_clients_connected, run_data_code_migrations, + unapplied_required_data_migrations, + }; + + struct Scratch { + db: PgPool, + admin: PgPool, + db_name: String, + } + + impl Scratch { + async fn cleanup(self) { + let Scratch { db, admin, db_name } = self; + db.close().await; + sqlx::query(&format!( + "DROP DATABASE IF EXISTS \"{db_name}\" WITH (FORCE)" + )) + .execute(&admin) + .await + .expect("drop the scratch database"); + admin.close().await; + } + } + + fn base_conn() -> Option { + let conn = std::env::var("EXTENDDB_TEST_PG_CONNECTION_STRING").ok()?; + (!conn.trim().is_empty()).then(|| conn.trim_end_matches('/').to_owned()) + } + + fn skip(test: &str) { + eprintln!( + "SKIP {test}: EXTENDDB_TEST_PG_CONNECTION_STRING is not set, so there is no \ + PostgreSQL to build a scratch database in." + ); + } + + async fn scratch() -> Scratch { + let base = base_conn().expect("caller checks base_conn() first"); + let db_name = format!("eddb_esc_{}", uuid::Uuid::new_v4().simple())[..24].to_owned(); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect(&format!("{base}/postgres")) + .await + .expect("connect to the postgres maintenance database"); + sqlx::query(&format!("CREATE DATABASE \"{db_name}\"")) + .execute(&admin) + .await + .expect("create the scratch database"); + // Named like the migrate command, so the connection guard on 004 treats + // this pool as the migration's own. + let db = PgPoolOptions::new() + .max_connections(2) + .connect_with( + format!("{base}/{db_name}") + .parse::() + .expect("scratch connection string") + .application_name(MIGRATE_APPLICATION_NAME), + ) + .await + .expect("connect to the scratch database"); + for (_, sql) in CATALOG_MIGRATIONS.iter().chain(DATA_MIGRATIONS) { + sqlx::raw_sql(sql) + .execute(&db) + .await + .expect("apply a shipped migration"); + } + // 003 creates its indexes from parsed catalog key schemas and is not + // what these tests exercise, so record it as already applied. + sqlx::query("INSERT INTO schema_history (filename) VALUES ('003_gsi_base_key_index')") + .execute(&db) + .await + .expect("record 003 as applied"); + Scratch { db, admin, db_name } + } + + /// Seed a data table, an index table, `gsi_pending`, `stream_records`, and + /// `backup_items` with three row kinds: raw U+0001 in key columns and + /// inside jsonb strings and object keys, the literal six-character text + /// backslash-u-0001 inside a jsonb string, and plain rows. + async fn seed(db: &PgPool) { + let ks = json!([ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "sk", "KeyType": "RANGE"} + ]); + let ad = json!([ + {"AttributeName": "pk", "AttributeType": "S"}, + {"AttributeName": "sk", "AttributeType": "S"} + ]); + sqlx::query("INSERT INTO accounts (account_id, account_name) VALUES ($1, $2)") + .bind("111122223333") + .bind("esc-check") + .execute(db) + .await + .expect("seed the account row"); + sqlx::query( + "INSERT INTO tables (account_id, table_name, key_schema, attribute_definitions, \ + table_status, table_arn, table_id) \ + VALUES ($1, 't', $2, $3, 'ACTIVE', \ + 'arn:aws:dynamodb:us-east-1:111122223333:table/t', 't1')", + ) + .bind("111122223333") + .bind(&ks) + .bind(&ad) + .execute(db) + .await + .expect("seed the tables row"); + sqlx::query( + "INSERT INTO indexes (table_id, index_id, index_name, index_type, key_schema, \ + projection) VALUES ('t1', 'i1', 'gsi1', 'GSI', $1, $2)", + ) + .bind(json!([{"AttributeName": "gpk", "KeyType": "HASH"}])) + .bind(json!({"ProjectionType": "ALL"})) + .execute(db) + .await + .expect("seed the indexes row"); + + sqlx::raw_sql( + r#"CREATE TABLE "_ddb_t1" ( + pk TEXT COLLATE "C" NOT NULL, + sk_s TEXT COLLATE "C", + sk_n NUMERIC, + sk_b BYTEA, + item_data JSONB NOT NULL, + PRIMARY KEY (pk, sk_s) + ); + CREATE TABLE "_ddb_i1" ( + pk TEXT COLLATE "C" NOT NULL, + sk_s TEXT COLLATE "C", + sk_n NUMERIC, + sk_b BYTEA, + base_pk TEXT COLLATE "C" NOT NULL, + base_sk_s TEXT COLLATE "C", + base_sk_n NUMERIC, + base_sk_b BYTEA, + item_data JSONB NOT NULL, + PRIMARY KEY (pk, base_pk, base_sk_s) + );"#, + ) + .execute(db) + .await + .expect("create the physical data and index tables"); + + for (pk, sk, item) in [ + ( + "a\u{1}b", + "s\u{1}", + json!({"pk": {"S": "a\u{1}b"}, "k\u{1}": {"S": "v\u{1}"}}), + ), + ("lit", "s", json!({"note": {"S": "x\\u0001y"}})), + ("plain", "s", json!({"pk": {"S": "plain"}})), + ] { + sqlx::query("INSERT INTO \"_ddb_t1\" (pk, sk_s, item_data) VALUES ($1, $2, $3)") + .bind(pk) + .bind(sk) + .bind(item) + .execute(db) + .await + .expect("seed a data table row"); + } + for (pk, sk, base_pk, base_sk, item) in [ + ( + "g\u{1}", + "gs\u{1}", + "a\u{1}b", + "s\u{1}", + json!({"gpk": {"S": "g\u{1}"}}), + ), + ("zz", "zs", "plain", "s", json!({"gpk": {"S": "zz"}})), + ] { + sqlx::query( + "INSERT INTO \"_ddb_i1\" (pk, sk_s, base_pk, base_sk_s, item_data) \ + VALUES ($1, $2, $3, $4, $5)", + ) + .bind(pk) + .bind(sk) + .bind(base_pk) + .bind(base_sk) + .bind(item) + .execute(db) + .await + .expect("seed an index table row"); + } + + for (new_item, context) in [ + (json!({"pk": {"S": "a\u{1}b"}}), json!({"k\u{1}": 1})), + (json!({"pk": {"S": "p"}}), json!({"k": 1})), + ] { + sqlx::query( + "INSERT INTO gsi_pending (table_id, worker_partition, old_item, new_item, \ + index_context) VALUES ('t1', 0, NULL, $1, $2)", + ) + .bind(new_item) + .bind(context) + .execute(db) + .await + .expect("seed a gsi_pending row"); + } + + sqlx::query( + "INSERT INTO stream_shards (shard_id, table_id, starting_sequence_number) \ + VALUES ('shard-1', 't1', '1')", + ) + .execute(db) + .await + .expect("seed the stream shard"); + for (seq, record) in [ + ("1", json!({"Keys": {"pk": {"S": "a\u{1}b"}}})), + ("2", json!({"Keys": {"pk": {"S": "p"}}})), + ] { + sqlx::query( + "INSERT INTO stream_records (shard_id, sequence_number, table_id, event_name, \ + record_data) VALUES ('shard-1', $1, 't1', 'INSERT', $2)", + ) + .bind(seq) + .bind(record) + .execute(db) + .await + .expect("seed a stream record"); + } + + sqlx::query( + "INSERT INTO backups (backup_arn, backup_name, table_id, table_name, account_id, \ + key_schema, attribute_definitions) \ + VALUES ('arn:aws:dynamodb:us-east-1:111122223333:table/t/backup/b1', 'b1', 't1', \ + 't', '111122223333', $1, $2)", + ) + .bind(&ks) + .bind(&ad) + .execute(db) + .await + .expect("seed the backups row"); + for (pk, sk, item) in [ + ("p\u{1}", Some("s\u{1}"), json!({"pk": {"S": "p\u{1}"}})), + ("plain", None, json!({"pk": {"S": "plain"}})), + ] { + sqlx::query( + "INSERT INTO backup_items (backup_arn, pk, sk, item_data) \ + VALUES ('arn:aws:dynamodb:us-east-1:111122223333:table/t/backup/b1', $1, $2, $3)", + ) + .bind(pk) + .bind(sk) + .bind(item) + .execute(db) + .await + .expect("seed a backup item"); + } + } + + #[tokio::test] + async fn reencodes_legacy_rows_once_and_records_itself() { + if base_conn().is_none() { + skip("reencodes_legacy_rows_once_and_records_itself"); + return; + } + let s = scratch().await; + seed(&s.db).await; + + let pending = pending_data_migrations(&s.db).await.expect("pending list"); + assert!( + pending.iter().any(|p| p == "004_escape_control_chars"), + "004 must be pending before the run: {pending:?}" + ); + + run_data_code_migrations(&s.db, &s.db) + .await + .expect("run the data code migrations"); + + // Data table: raw U+0001 became U+0001 U+0002 in the key columns and + // inside every jsonb string and object key; the literal-backslash row + // and the plain row are byte-identical. + let rows: Vec<(String, String, serde_json::Value)> = + sqlx::query_as("SELECT pk, sk_s, item_data FROM \"_ddb_t1\" ORDER BY pk") + .fetch_all(&s.db) + .await + .expect("read the data table"); + assert_eq!( + rows, + vec![ + ( + "a\u{1}\u{2}b".to_owned(), + "s\u{1}\u{2}".to_owned(), + json!({"pk": {"S": "a\u{1}\u{2}b"}, "k\u{1}\u{2}": {"S": "v\u{1}\u{2}"}}), + ), + ( + "lit".to_owned(), + "s".to_owned(), + json!({"note": {"S": "x\\u0001y"}}), + ), + ( + "plain".to_owned(), + "s".to_owned(), + json!({"pk": {"S": "plain"}}), + ), + ] + ); + + // Index table: pk, sk_s, base_pk, base_sk_s and item_data all re-encoded. + let rows: Vec<(String, String, String, String, serde_json::Value)> = sqlx::query_as( + "SELECT pk, sk_s, base_pk, base_sk_s, item_data FROM \"_ddb_i1\" ORDER BY pk", + ) + .fetch_all(&s.db) + .await + .expect("read the index table"); + assert_eq!( + rows, + vec![ + ( + "g\u{1}\u{2}".to_owned(), + "gs\u{1}\u{2}".to_owned(), + "a\u{1}\u{2}b".to_owned(), + "s\u{1}\u{2}".to_owned(), + json!({"gpk": {"S": "g\u{1}\u{2}"}}), + ), + ( + "zz".to_owned(), + "zs".to_owned(), + "plain".to_owned(), + "s".to_owned(), + json!({"gpk": {"S": "zz"}}), + ), + ] + ); + + // gsi_pending: NULL old_item stays NULL, the raw row is re-encoded in + // both jsonb columns, the plain row is untouched. + let rows: Vec<( + Option, + serde_json::Value, + serde_json::Value, + )> = + sqlx::query_as("SELECT old_item, new_item, index_context FROM gsi_pending ORDER BY id") + .fetch_all(&s.db) + .await + .expect("read gsi_pending"); + assert_eq!( + rows, + vec![ + ( + None, + json!({"pk": {"S": "a\u{1}\u{2}b"}}), + json!({"k\u{1}\u{2}": 1}), + ), + (None, json!({"pk": {"S": "p"}}), json!({"k": 1})), + ] + ); + + // stream_records. + let rows: Vec<(serde_json::Value,)> = + sqlx::query_as("SELECT record_data FROM stream_records ORDER BY sequence_number") + .fetch_all(&s.db) + .await + .expect("read stream_records"); + assert_eq!( + rows, + vec![ + (json!({"Keys": {"pk": {"S": "a\u{1}\u{2}b"}}}),), + (json!({"Keys": {"pk": {"S": "p"}}}),), + ] + ); + + // backup_items, including the NULL sk row. + let rows: Vec<(String, Option, serde_json::Value)> = + sqlx::query_as("SELECT pk, sk, item_data FROM backup_items ORDER BY pk") + .fetch_all(&s.db) + .await + .expect("read backup_items"); + assert_eq!( + rows, + vec![ + ( + "p\u{1}\u{2}".to_owned(), + Some("s\u{1}\u{2}".to_owned()), + json!({"pk": {"S": "p\u{1}\u{2}"}}), + ), + ("plain".to_owned(), None, json!({"pk": {"S": "plain"}}),), + ] + ); + + // Recorded in schema_history, one progress marker per table, and no + // longer pending. + let recorded: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM schema_history \ + WHERE filename = '004_escape_control_chars')", + ) + .fetch_one(&s.db) + .await + .expect("read schema_history"); + assert!(recorded, "schema_history must record the migration"); + let markers: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM data_code_migration_progress \ + WHERE migration_name = '004_escape_control_chars'", + ) + .fetch_one(&s.db) + .await + .expect("count progress markers"); + assert_eq!(markers, 5, "one marker per processed table"); + let pending = pending_data_migrations(&s.db).await.expect("pending list"); + assert!( + !pending.iter().any(|p| p == "004_escape_control_chars"), + "004 must not be pending after the run: {pending:?}" + ); + + // A second run is skipped by the is_migration_applied guard and does + // not double-escape anything. + run_data_code_migrations(&s.db, &s.db) + .await + .expect("run the data code migrations again"); + let (pk, item): (String, serde_json::Value) = + sqlx::query_as("SELECT pk, item_data FROM \"_ddb_t1\" WHERE sk_s = $1") + .bind("s\u{1}\u{2}") + .fetch_one(&s.db) + .await + .expect("re-read the re-encoded row"); + assert_eq!(pk, "a\u{1}\u{2}b"); + assert_eq!( + item, + json!({"pk": {"S": "a\u{1}\u{2}b"}, "k\u{1}\u{2}": {"S": "v\u{1}\u{2}"}}) + ); + + s.cleanup().await; + } + + /// The jsonb candidate predicate must not depend on the session's + /// standard_conforming_strings: with it off, a plain '\u0001' literal is a + /// single 0x01 byte that never appears in jsonb text and matches nothing, + /// so every jsonb row would be skipped while the migration recorded itself. + #[tokio::test] + async fn reencodes_jsonb_rows_with_standard_conforming_strings_off() { + if base_conn().is_none() { + skip("reencodes_jsonb_rows_with_standard_conforming_strings_off"); + return; + } + let s = scratch().await; + seed(&s.db).await; + + let base = base_conn().unwrap(); + let scs_off = PgPoolOptions::new() + .max_connections(2) + .after_connect(|conn, _| { + Box::pin(async move { + sqlx::Executor::execute(&mut *conn, "SET standard_conforming_strings = off") + .await?; + Ok(()) + }) + }) + .connect_with( + format!("{base}/{}", s.db_name) + .parse::() + .unwrap() + .application_name(MIGRATE_APPLICATION_NAME), + ) + .await + .expect("connect with standard_conforming_strings off"); + + run_data_code_migrations(&scs_off, &scs_off) + .await + .expect("run 004"); + + let (pk, sk, item): (String, String, serde_json::Value) = + sqlx::query_as("SELECT pk, sk_s, item_data FROM \"_ddb_t1\" WHERE pk LIKE 'a%'") + .fetch_one(&s.db) + .await + .unwrap(); + assert_eq!(pk, "a\u{1}\u{2}b"); + assert_eq!(sk, "s\u{1}\u{2}"); + assert_eq!( + item, + json!({"pk": {"S": "a\u{1}\u{2}b"}, "k\u{1}\u{2}": {"S": "v\u{1}\u{2}"}}), + "jsonb rows must be re-encoded whatever the session setting" + ); + scs_off.close().await; + s.cleanup().await; + } + + /// The server must not serve a data database that has not been through + /// 004: rows it wrote through the escape would be rewritten by the + /// migration later. The required list is empty only once 004 is recorded. + #[tokio::test] + async fn required_migrations_are_reported_until_004_is_recorded() { + if base_conn().is_none() { + skip("required_migrations_are_reported_until_004_is_recorded"); + return; + } + let s = scratch().await; + seed(&s.db).await; + + let missing = unapplied_required_data_migrations(&s.db).await.unwrap(); + assert_eq!(missing, vec!["004_escape_control_chars".to_owned()]); + + run_data_code_migrations(&s.db, &s.db) + .await + .expect("run 004"); + + let missing = unapplied_required_data_migrations(&s.db).await.unwrap(); + assert!(missing.is_empty(), "{missing:?}"); + s.cleanup().await; + } + + /// A server (or any other session) still connected to the data database + /// blocks 004; the migrate command's own connections do not. + #[tokio::test] + async fn refuses_004_while_another_client_is_connected() { + if base_conn().is_none() { + skip("refuses_004_while_another_client_is_connected"); + return; + } + let s = scratch().await; + seed(&s.db).await; + let base = base_conn().unwrap(); + + // An unnamed pool stands in for a running server. + let server_pool = PgPoolOptions::new() + .max_connections(1) + .connect(&format!("{base}/{}", s.db_name)) + .await + .expect("connect as a server would"); + sqlx::query("SELECT 1").execute(&server_pool).await.unwrap(); + + let err = refuse_if_other_clients_connected(&s.db) + .await + .expect_err("an unnamed connection must block the migration"); + let text = format!("{err:?}"); + assert!( + text.contains("other connection(s) hold the data database"), + "{text}" + ); + assert!(text.contains(IGNORE_CONNECTIONS_ENV), "{text}"); + let err = run_data_code_migrations(&s.db, &s.db) + .await + .expect_err("004 must not run while the server is connected"); + assert!(format!("{err:?}").contains("other connection(s)")); + assert!( + unapplied_required_data_migrations(&s.db) + .await + .unwrap() + .len() + == 1, + "nothing may be recorded when the guard refuses" + ); + + // Server stopped: only the migrate command's own connections remain. + server_pool.close().await; + refuse_if_other_clients_connected(&s.db) + .await + .expect("no foreign connection left"); + run_data_code_migrations(&s.db, &s.db) + .await + .expect("run 004"); + s.cleanup().await; + } +} diff --git a/crates/storage-postgres/src/stream_engine.rs b/crates/storage-postgres/src/stream_engine.rs index 151a90180..059b26ce2 100755 --- a/crates/storage-postgres/src/stream_engine.rs +++ b/crates/storage-postgres/src/stream_engine.rs @@ -94,8 +94,7 @@ impl StreamEngine for PostgresEngine { let shard_id = shard_id.to_string(); let table_name = table_name.to_string(); Box::pin(async move { - let record_json = - serde_json::to_value(&record).map_err(|e| StorageError::Internal(e.to_string()))?; + let record_json = crate::data::to_stored_json(&record)?; let table_id: String = sqlx::query_scalar( "SELECT table_id FROM tables WHERE account_id = $1 AND table_name = $2", @@ -197,9 +196,7 @@ impl StreamEngine for PostgresEngine { let records: Vec = rows .into_iter() - .map(|(data,)| { - serde_json::from_value(data).map_err(|e| StorageError::Internal(e.to_string())) - }) + .map(|(data,)| crate::data::stored_json_to(data)) .collect::, _>>()?; let last_seq = records.last().map(|r| r.dynamodb.sequence_number.clone()); diff --git a/crates/storage-postgres/src/vector_search.rs b/crates/storage-postgres/src/vector_search.rs index 5b278d467..5f4106acf 100644 --- a/crates/storage-postgres/src/vector_search.rs +++ b/crates/storage-postgres/src/vector_search.rs @@ -277,10 +277,17 @@ impl VectorSearchEngine for PostgresEngine { .bind(partition.into_bytes()) .bind(top_k) .bind(query_norm); + // The stored item carries the same escape as every other item_data + // column, so the filter's attribute name and value must be compared + // in stored form. for (name, value) in &filters { - let value_json = serde_json::to_string(value) - .map_err(|e| StorageError::Internal(format!("filter value: {e}")))?; - query = query.bind(name).bind(value_json); + let value_json = serde_json::to_value(value) + .map(extenddb_storage::util::escape_json_strings) + .map_err(|e| StorageError::Internal(format!("filter value: {e}")))? + .to_string(); + query = query + .bind(extenddb_storage::util::escape_control(name).into_owned()) + .bind(value_json); } let rows = query.fetch_all(&self.data_pool).await.map_err(|e| { @@ -298,7 +305,7 @@ impl VectorSearchEngine for PostgresEngine { let mut hits = Vec::with_capacity(rows.len()); for (ordered, embedding, item_json) in rows { - let mut item: extenddb_core::types::Item = serde_json::from_value(item_json) + let mut item: extenddb_core::types::Item = crate::data::json_to_item(item_json) .map_err(|e| StorageError::Internal(format!("stored item: {e}")))?; // Reinstated from the stored f32s rather than from a second copy in // the payload, so what comes back is the narrowed value that was diff --git a/crates/storage/src/util/control_chars.rs b/crates/storage/src/util/control_chars.rs new file mode 100644 index 000000000..07db91310 --- /dev/null +++ b/crates/storage/src/util/control_chars.rs @@ -0,0 +1,332 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Order-preserving escape for the two characters a storage engine may not be +//! able to hold verbatim. +//! +//! DynamoDB accepts the character U+0000 anywhere a string appears: partition +//! and sort keys, index keys, attribute values, map keys, attribute names. Two +//! of the storage engines cannot store it as is. PostgreSQL `TEXT` rejects the +//! byte (`invalid byte sequence for encoding "UTF8": 0x00`) and `jsonb` rejects +//! the `\u0000` escape. BSON field names are C strings and end at the first NUL. +//! +//! The escape maps each affected character to a two-character sequence that +//! contains no NUL: +//! +//! | input | output | +//! |--------|---------------| +//! | U+0000 | U+0001 U+0001 | +//! | U+0001 | U+0001 U+0002 | +//! +//! Every other character passes through unchanged, so a string without either +//! character encodes to itself and the stored form of ordinary data does not +//! change. The mapping is applied per character and the two output sequences +//! share the prefix U+0001, so three properties hold, each pinned by a test +//! below: +//! +//! - **Round trip.** `unescape(escape(s)) == s` for every `s`. +//! - **Byte order.** `escape(a) < escape(b)` exactly when `a < b` under UTF-8 +//! byte comparison. PostgreSQL compares key columns under `COLLATE "C"`, so +//! Query ordering, `BETWEEN`, and cursor comparisons on escaped keys give the +//! same answers as on the originals. +//! - **Prefix.** `escape(p)` is a prefix of `escape(s)` exactly when `p` is a +//! prefix of `s`, so `begins_with` on escaped values is exact. +//! +//! Decoding is lenient about a U+0001 that is not followed by U+0001 or U+0002: +//! it is returned as a literal U+0001. Rows written before this escape existed +//! can hold such a character, and a lenient decoder reads them unchanged. A +//! legacy string that contained U+0001 immediately followed by U+0001 or U+0002 +//! would decode differently; the operator migration `004_escape_control_chars` +//! re-encodes such rows so that case cannot arise after it has run. + +use std::borrow::Cow; + +const ESC: char = '\u{1}'; +const NUL_TAIL: char = '\u{1}'; +const ESC_TAIL: char = '\u{2}'; + +/// Whether `s` contains a character the escape changes. +#[must_use] +pub fn needs_escape(s: &str) -> bool { + s.bytes().any(|b| b == 0 || b == 1) +} + +/// Escape U+0000 and U+0001 in `s`. Borrows when there is nothing to change. +#[must_use] +pub fn escape_control(s: &str) -> Cow<'_, str> { + if !needs_escape(s) { + return Cow::Borrowed(s); + } + let mut out = String::with_capacity(s.len() + 8); + for ch in s.chars() { + match ch { + '\u{0}' => { + out.push(ESC); + out.push(NUL_TAIL); + } + '\u{1}' => { + out.push(ESC); + out.push(ESC_TAIL); + } + other => out.push(other), + } + } + Cow::Owned(out) +} + +/// Reverse [`escape_control`]. Borrows when there is nothing to change. A +/// U+0001 not followed by U+0001 or U+0002 is kept as a literal U+0001 (see +/// the module documentation for why). +#[must_use] +pub fn unescape_control(s: &str) -> Cow<'_, str> { + if !s.bytes().any(|b| b == 1) { + return Cow::Borrowed(s); + } + let mut out = String::with_capacity(s.len()); + let mut chars = s.chars().peekable(); + while let Some(ch) = chars.next() { + if ch == ESC { + match chars.peek() { + Some(&NUL_TAIL) => { + chars.next(); + out.push('\u{0}'); + } + Some(&ESC_TAIL) => { + chars.next(); + out.push('\u{1}'); + } + _ => out.push('\u{1}'), + } + } else { + out.push(ch); + } + } + Cow::Owned(out) +} + +/// Apply [`escape_control`] to every string in a JSON tree: object keys and +/// string values, at every depth. Used where the whole document lands in a +/// store that rejects U+0000 in any string (PostgreSQL `jsonb`). +#[must_use] +pub fn escape_json_strings(v: serde_json::Value) -> serde_json::Value { + map_json(v, true, &|s| escape_control(s).into_owned()) +} + +/// Reverse [`escape_json_strings`]. +#[must_use] +pub fn unescape_json_strings(v: serde_json::Value) -> serde_json::Value { + map_json(v, true, &|s| unescape_control(s).into_owned()) +} + +/// Apply [`escape_control`] to every object key in a JSON tree, leaving string +/// values alone. Used where only field names are restricted (BSON). +#[must_use] +pub fn escape_json_keys(v: serde_json::Value) -> serde_json::Value { + map_json(v, false, &|s| escape_control(s).into_owned()) +} + +/// Reverse [`escape_json_keys`]. +#[must_use] +pub fn unescape_json_keys(v: serde_json::Value) -> serde_json::Value { + map_json(v, false, &|s| unescape_control(s).into_owned()) +} + +fn map_json( + v: serde_json::Value, + values_too: bool, + f: &dyn Fn(&str) -> String, +) -> serde_json::Value { + use serde_json::Value; + match v { + Value::String(s) if values_too => { + if needs_escape(&s) { + Value::String(f(&s)) + } else { + Value::String(s) + } + } + Value::Array(items) => Value::Array( + items + .into_iter() + .map(|item| map_json(item, values_too, f)) + .collect(), + ), + Value::Object(map) => { + let mut out = serde_json::Map::with_capacity(map.len()); + for (k, val) in map { + let key = if k.bytes().any(|b| b == 0 || b == 1) { + f(&k) + } else { + k + }; + out.insert(key, map_json(val, values_too, f)); + } + Value::Object(out) + } + other => other, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + /// Every string over the alphabet {NUL, U+0001, U+0002, 'a'} up to the + /// given length. Small alphabet, exhaustive: the escape only looks at + /// these three code points and everything else is identity. + fn corpus(max_len: usize) -> Vec { + let alphabet = ['\u{0}', '\u{1}', '\u{2}', 'a']; + let mut out = vec![String::new()]; + let mut frontier = vec![String::new()]; + for _ in 0..max_len { + let mut next = Vec::new(); + for s in &frontier { + for ch in alphabet { + let mut t = s.clone(); + t.push(ch); + next.push(t); + } + } + out.extend(next.iter().cloned()); + frontier = next; + } + out + } + + #[test] + fn plain_strings_are_borrowed_unchanged() { + for s in ["", "a", "hello", "\u{2}\u{3}", "\u{e9}\u{4e2d}", "a\tb\nc"] { + assert!(matches!(escape_control(s), Cow::Borrowed(_)), "{s:?}"); + assert!(matches!(unescape_control(s), Cow::Borrowed(_)), "{s:?}"); + assert_eq!(escape_control(s), s); + assert_eq!(unescape_control(s), s); + } + } + + #[test] + fn table_of_the_two_mappings() { + assert_eq!(escape_control("\u{0}"), "\u{1}\u{1}"); + assert_eq!(escape_control("\u{1}"), "\u{1}\u{2}"); + assert_eq!(escape_control("a\u{0}b"), "a\u{1}\u{1}b"); + assert_eq!(escape_control("\u{0}\u{1}"), "\u{1}\u{1}\u{1}\u{2}"); + assert_eq!(unescape_control("\u{1}\u{1}"), "\u{0}"); + assert_eq!(unescape_control("\u{1}\u{2}"), "\u{1}"); + } + + #[test] + fn escaped_output_never_contains_nul() { + for s in corpus(5) { + assert!(!escape_control(&s).contains('\u{0}'), "{s:?}"); + } + } + + #[test] + fn round_trips_every_string_in_the_corpus() { + for s in corpus(5) { + assert_eq!(unescape_control(&escape_control(&s)), s, "{s:?}"); + } + for s in [ + "caf\u{e9}\u{0}\u{4e2d}", + "\u{0}\u{0}\u{0}", + "\u{1}\u{1}\u{1}", + "x\u{1}\u{0}\u{2}y", + ] { + assert_eq!(unescape_control(&escape_control(s)), s, "{s:?}"); + } + } + + #[test] + fn escape_preserves_utf8_byte_order() { + // The whole point: PostgreSQL compares the escaped column under + // COLLATE "C", and every ordering answer must equal the answer on + // the originals. + let corpus = corpus(4); + for a in &corpus { + for b in &corpus { + let want = a.as_bytes().cmp(b.as_bytes()); + let got = escape_control(a) + .as_bytes() + .cmp(escape_control(b).as_bytes()); + assert_eq!(want, got, "a={a:?} b={b:?}"); + } + } + } + + #[test] + fn escape_preserves_prefixes() { + // begins_with on escaped values must be exact in both directions. + let corpus = corpus(4); + for p in &corpus { + for s in &corpus { + let want = s.as_bytes().starts_with(p.as_bytes()); + let got = escape_control(s) + .as_bytes() + .starts_with(escape_control(p).as_bytes()); + assert_eq!(want, got, "p={p:?} s={s:?}"); + } + } + } + + #[test] + fn escaped_forms_are_distinct_for_distinct_inputs() { + let corpus = corpus(4); + let mut seen = std::collections::HashMap::new(); + for s in &corpus { + let e = escape_control(s).into_owned(); + if let Some(prev) = seen.insert(e.clone(), s.clone()) { + panic!("collision: {prev:?} and {s:?} both escape to {e:?}"); + } + } + } + + #[test] + fn lenient_decode_keeps_a_stray_escape_character_literal() { + // A row written before the escape existed can carry a raw U+0001 + // followed by anything other than U+0001 or U+0002. + assert_eq!(unescape_control("a\u{1}b"), "a\u{1}b"); + assert_eq!(unescape_control("\u{1}"), "\u{1}"); + assert_eq!(unescape_control("\u{1}\u{3}"), "\u{1}\u{3}"); + assert_eq!(unescape_control("x\u{1}"), "x\u{1}"); + } + + #[test] + fn json_strings_variant_touches_keys_and_values_at_every_depth() { + let item = json!({ + "a\u{0}b": {"S": "v\u{0}"}, + "m": {"M": {"k\u{1}": {"S": "\u{0}"}, "plain": {"N": "1"}}}, + "l": {"L": [{"S": "\u{0}x"}, {"SS": ["\u{0}", "y"]}]}, + "b": {"B": "AAEC"} + }); + let escaped = escape_json_strings(item.clone()); + assert!(!escaped.to_string().contains("\\u0000"), "{escaped}"); + assert_eq!(escaped["a\u{1}\u{1}b"]["S"], json!("v\u{1}\u{1}")); + assert_eq!(escaped["m"]["M"]["k\u{1}\u{2}"]["S"], json!("\u{1}\u{1}")); + assert_eq!(escaped["l"]["L"][1]["SS"][0], json!("\u{1}\u{1}")); + assert_eq!(escaped["b"]["B"], json!("AAEC")); + assert_eq!(unescape_json_strings(escaped), item); + } + + #[test] + fn json_keys_variant_leaves_values_alone() { + let item = json!({ + "a\u{0}b": {"S": "v\u{0}"}, + "m": {"M": {"k\u{0}": {"L": [{"S": "\u{0}"}]}}} + }); + let escaped = escape_json_keys(item.clone()); + assert_eq!(escaped["a\u{1}\u{1}b"]["S"], json!("v\u{0}")); + assert_eq!( + escaped["m"]["M"]["k\u{1}\u{1}"]["L"][0]["S"], + json!("\u{0}") + ); + assert_eq!(unescape_json_keys(escaped), item); + } + + #[test] + fn json_without_control_characters_is_unchanged() { + let item = json!({"pk": {"S": "p"}, "n": {"N": "1"}, "m": {"M": {"k": {"S": "v"}}}}); + assert_eq!(escape_json_strings(item.clone()), item); + assert_eq!(escape_json_keys(item.clone()), item); + assert_eq!(unescape_json_strings(item.clone()), item); + assert_eq!(unescape_json_keys(item.clone()), item); + } +} diff --git a/crates/storage/src/util/key.rs b/crates/storage/src/util/key.rs index 1519afd9f..4453b794b 100755 --- a/crates/storage/src/util/key.rs +++ b/crates/storage/src/util/key.rs @@ -23,9 +23,10 @@ pub enum SortKeyValue { /// /// For single-attribute keys, returns the value directly (no encoding). /// For multi-attribute keys, uses netstring encoding: each part is encoded as -/// `:,` and concatenated. This is provably collision-free -/// regardless of value content, and compatible with `PostgreSQL` TEXT columns -/// (no null bytes). +/// `:,` and concatenated. This is collision-free +/// regardless of value content. The parts are the raw attribute strings; a +/// backend whose columns cannot hold every character (PostgreSQL, see +/// `escape_control`) escapes its key text in its own wrapper. pub fn composite_pk_to_text( item: &Item, key_schema: &[KeySchemaElement], @@ -286,7 +287,7 @@ pub fn recover_sort_key_definitions( /// Encode multiple string parts into a single netstring-encoded composite key. /// /// Format: `:,:,...` — e.g., `"abc"` + `"de"` → `"3:abc,2:de,"`. -/// This encoding is unambiguous for arbitrary byte content and contains no null bytes. +/// This encoding is unambiguous for arbitrary part content; the parts are copied as given. #[must_use] pub fn encode_netstring_composite(parts: &[String]) -> String { let mut out = String::new(); diff --git a/crates/storage/src/util/mod.rs b/crates/storage/src/util/mod.rs index 81e2ae850..3f1d08a06 100755 --- a/crates/storage/src/util/mod.rs +++ b/crates/storage/src/util/mod.rs @@ -7,9 +7,14 @@ //! id validation. mod arn; +mod control_chars; mod key; pub use arn::{index_arn, parse_stream_arn, stream_arn, table_arn}; +pub use control_chars::{ + escape_control, escape_json_keys, escape_json_strings, needs_escape, unescape_control, + unescape_json_keys, unescape_json_strings, +}; pub use key::SortKeyValue; pub use key::{ composite_pk_to_text, effective_attribute_definitions, encode_netstring_composite, diff --git a/docs/design/04-component-storage.md b/docs/design/04-component-storage.md index b85c60232..5b6be7fb5 100755 --- a/docs/design/04-component-storage.md +++ b/docs/design/04-component-storage.md @@ -402,6 +402,25 @@ The PostgreSQL backend uses two categories of tables: attributes, matching the DynamoDB model where key attributes are part of the item. +- **U+0000 in strings**: DynamoDB accepts the character U+0000 anywhere a + string appears (measured 2026-09-18: partition and sort keys, index keys, + values, strings in lists, map keys, attribute names; it sorts as the byte + 0x00 and is a different key from the six-character text `\u0000`). + PostgreSQL `TEXT` rejects the byte and `jsonb` rejects the `\u0000` escape, + so every string that reaches a key column or an `item_data`, `gsi_pending`, + `stream_records`, or `backup_items` document goes through the escape in + `extenddb_storage::util::escape_control`: U+0000 is stored as U+0001 U+0001 + and U+0001 as U+0001 U+0002; every other character is unchanged, so ordinary + data is stored exactly as before. The mapping preserves UTF-8 byte order and + prefixes, so `COLLATE "C"` comparisons, `BETWEEN`, `begins_with`, and + row-comparison cursors on escaped columns return the rows the raw text + would. SQL that addresses an attribute by name inside `item_data` (the TTL + index and sweep, vector search filters) uses the escaped name. The MongoDB + backend applies the same escape to BSON field names only, since BSON string + values are length-prefixed and hold the byte as is. Rows written before the + escape existed are re-encoded once by the operator migration + `004_escape_control_chars` (see the upgrade manual). + - **GSI tables**: GSI tables include base table primary key columns (`base_pk`, `base_sk_*`) as actual SQL columns (not just inside `item_data` JSONB). This is required because: (1) GSI keys are not unique — two base table items can diff --git a/docs/manuals/07-upgrade-manual.md b/docs/manuals/07-upgrade-manual.md index fe2bf5eb2..3445f26aa 100755 --- a/docs/manuals/07-upgrade-manual.md +++ b/docs/manuals/07-upgrade-manual.md @@ -212,6 +212,66 @@ The first release, so all 0.0.2 deployments were fresh installs. --- +## Data Migrations + +Some changes rewrite rows in the data database rather than the catalog schema. These are +implemented in Rust (`DATA_CODE_MIGRATIONS` in `crates/storage-postgres/src/migrations.rs`), +tracked in the data database's own `schema_history` table, and applied once by +`extenddb migrate --yes`. Running `extenddb migrate` without `--yes` lists them as pending +and changes nothing. + +### 004_escape_control_chars + +Strings are stored through an order-preserving escape so that U+0000, which DynamoDB +accepts anywhere a string appears, can be held in PostgreSQL `TEXT` and `jsonb` columns: +U+0000 is stored as U+0001 U+0001, and U+0001 as U+0001 U+0002. Rows written before the +escape existed hold their strings raw, and a raw U+0001 in such a row could be misread by +the decoder. This migration re-encodes those rows. + +What it rewrites: every data table and index table (the text key columns and the +`item_data` document), vector index tables, `gsi_pending`, `stream_records`, and +`backup_items`. Only rows that contain U+0001 are touched; a database with none is left as +it is, apart from the `schema_history` row. + +The migration scans every data and index table once (each row's `jsonb` document is +rendered to text to find candidates), so on a large deployment it takes time proportional +to the size of the data. It processes one table per transaction and keeps a per-table +progress marker in `data_code_migration_progress`, so an interrupted run resumes with the +tables that are not yet done and never rewrites a table twice. + +Order of operations. A server built with the escape refuses to start on a data database +that does not record this migration, with a message naming it, because rows the new build +writes are indistinguishable from legacy rows to the migration (both can contain U+0001) and +would be rewritten by it. Upgrade a PostgreSQL deployment in this order: + +1. Stop every ExtendDB server that uses the data database. A server from an earlier release + still writing during the migration would leave unescaped rows behind. +2. Run `extenddb migrate --yes` with the new binary. +3. Start the new servers. + +A fresh `extenddb init` records the migration, so new deployments start without this step. +A rolling deployment that mixes the two releases against one data database is not +supported for this release: the new server will not start until the migration has run, +and the old server must be stopped before it runs. + +`extenddb migrate --yes` refuses to run this migration while any other session is connected +to the data database, and lists them, because a server still writing during the rewrite +would leave rows in the wrong encoding. If the only other connections are idle ones held +by a connection pooler, set `EXTENDDB_MIGRATE_IGNORE_CONNECTIONS=1` to skip that check. +Nothing stops a server from an earlier release from being started against a migrated +data database afterwards; it would read escaped rows as literal data. Do not downgrade +past this release without restoring the data database from a backup taken before the +migration. + +MongoDB. The MongoDB backend applies the same escape to BSON field names (attribute +names and map keys) and has no migration and no startup check. Field names written by +an earlier release are stored raw. A name containing a lone U+0001 reads back unchanged, +as do all string values and all key strings; a name containing U+0001 immediately +followed by U+0001 or U+0002 reads back with that pair decoded as U+0000 or U+0001. +Deployments whose attribute names never contain U+0001 are unaffected. + +--- + ## Behavior Changes by Release Catalog upgrades change the schema; behavior changes alter how the running server diff --git a/tests/test_nul_strings.py b/tests/test_nul_strings.py new file mode 100644 index 000000000..0e1d8c511 --- /dev/null +++ b/tests/test_nul_strings.py @@ -0,0 +1,378 @@ +# SPDX-License-Identifier: Apache-2.0 +"""U+0000 and U+0001 inside strings: keys, index keys, values, and names. + +DynamoDB places no character restriction on strings. Measured 2026-09-18 +against the service on a scratch table: PutItem accepts the NUL character +inside a partition key, a sort key, a GSI key, a non-key value, a string in a +list, a map key, and a top-level attribute name, and GetItem returns every one +byte-identical. NUL sorts as the byte 0x00, before every other character; it is +a different character from the six-character text ``\\u0000``; range and prefix +conditions treat it as an ordinary byte; the GSI holds the item; DeleteItem on +the key removes it. These tests pin exactly that, so every backend has to store +the character rather than its escape text and has to keep it out of the way of +its own storage encoding. + +U+0001 is included because a storage encoding that handles NUL has to leave the +next code point alone as well. +""" + +from __future__ import annotations + +import time + +import pytest +from botocore.exceptions import ClientError + +from conftest import scoped_table, wait_for_gsi_items + +NUL = "\u0000" +SOH = "\u0001" +ESCAPE_TEXT = "\\u0000" # backslash, u, 0, 0, 0, 0: six characters, not a NUL + +PK = f"a{NUL}b" +PK_WITHOUT_NUL = "ab" +PK_SOH = f"a{SOH}b" + +# One partition, five sort keys whose byte order is: NUL, U+0001, space, the +# escape text (starts with 0x5c), then "a" (0x61). +SORT_KEYS = [NUL, SOH, " ", ESCAPE_TEXT, "a"] + +# The item under the NUL sort key carries the character in every kind of +# position the item model has. +RICH_ITEM = { + "pk": {"S": PK}, + "sk": {"S": NUL}, + "g": {"S": f"g{NUL}"}, + "v": {"S": f"v{NUL}"}, + f"n{NUL}n": {"S": "x"}, + "m": {"M": {f"k{NUL}": {"S": f"{NUL}"}, f"k{SOH}{SOH}": {"N": "1"}}}, + "l": {"L": [{"S": NUL}, {"S": f"{SOH}{NUL}{SOH}"}]}, + "ss": {"SS": [NUL, "a"]}, +} + + +def _query(client, table, **kwargs): + return client.query( + TableName=table, + KeyConditionExpression=kwargs.pop("cond", "pk = :p"), + ExpressionAttributeValues={":p": {"S": PK}, **kwargs.pop("values", {})}, + **kwargs, + ) + + +def _sks(resp): + return [i["sk"]["S"] for i in resp["Items"]] + + +@pytest.fixture(scope="class") +def nul_table(dynamodb_client): + """(S, S) table with GSI gsi1 on g; the partition PK holds SORT_KEYS, plus + a control item in the partition without the NUL and one in a partition + keyed with U+0001.""" + with scoped_table( + dynamodb_client, + attribute_definitions=[ + {"AttributeName": "pk", "AttributeType": "S"}, + {"AttributeName": "sk", "AttributeType": "S"}, + {"AttributeName": "g", "AttributeType": "S"}, + ], + key_schema=[ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "sk", "KeyType": "RANGE"}, + ], + GlobalSecondaryIndexes=[ + { + "IndexName": "gsi1", + "KeySchema": [{"AttributeName": "g", "KeyType": "HASH"}], + "Projection": {"ProjectionType": "ALL"}, + } + ], + ) as name: + dynamodb_client.put_item(TableName=name, Item=RICH_ITEM) + for sk in SORT_KEYS[1:]: + dynamodb_client.put_item( + TableName=name, Item={"pk": {"S": PK}, "sk": {"S": sk}, "g": {"S": "plain"}} + ) + dynamodb_client.put_item( + TableName=name, Item={"pk": {"S": PK_WITHOUT_NUL}, "sk": {"S": "x"}, "g": {"S": "plain"}} + ) + dynamodb_client.put_item( + TableName=name, + Item={"pk": {"S": PK_SOH}, "sk": {"S": SOH}, "g": {"S": f"g{SOH}"}, "v": {"S": SOH}}, + ) + yield name + + +class TestNulInStrings: + def test_item_with_nul_everywhere_round_trips(self, dynamodb_client, nul_table): + resp = dynamodb_client.get_item( + TableName=nul_table, Key={"pk": {"S": PK}, "sk": {"S": NUL}}, ConsistentRead=True + ) + assert resp["Item"] == RICH_ITEM + + def test_soh_item_round_trips(self, dynamodb_client, nul_table): + resp = dynamodb_client.get_item( + TableName=nul_table, Key={"pk": {"S": PK_SOH}, "sk": {"S": SOH}}, ConsistentRead=True + ) + assert resp["Item"] == { + "pk": {"S": PK_SOH}, + "sk": {"S": SOH}, + "g": {"S": f"g{SOH}"}, + "v": {"S": SOH}, + } + + def test_partition_order_is_byte_order_with_nul_first(self, dynamodb_client, nul_table): + assert _sks(_query(dynamodb_client, nul_table, ConsistentRead=True)) == SORT_KEYS + assert ( + _sks(_query(dynamodb_client, nul_table, ScanIndexForward=False, ConsistentRead=True)) + == SORT_KEYS[::-1] + ) + + def test_nul_and_its_escape_text_are_different_keys(self, dynamodb_client, nul_table): + by_nul = dynamodb_client.get_item( + TableName=nul_table, Key={"pk": {"S": PK}, "sk": {"S": NUL}}, ConsistentRead=True + )["Item"] + by_text = dynamodb_client.get_item( + TableName=nul_table, Key={"pk": {"S": PK}, "sk": {"S": ESCAPE_TEXT}}, ConsistentRead=True + )["Item"] + assert by_nul["sk"]["S"] == NUL and len(by_nul["sk"]["S"]) == 1 + assert by_text["sk"]["S"] == ESCAPE_TEXT and len(by_text["sk"]["S"]) == 6 + + def test_partition_key_with_nul_is_distinct_from_the_key_without_it( + self, dynamodb_client, nul_table + ): + resp = dynamodb_client.get_item( + TableName=nul_table, Key={"pk": {"S": PK_WITHOUT_NUL}, "sk": {"S": NUL}}, ConsistentRead=True + ) + assert "Item" not in resp + resp = _query(dynamodb_client, nul_table, values={":p": {"S": PK_WITHOUT_NUL}}, ConsistentRead=True) + assert _sks(resp) == ["x"] + + def test_range_conditions_treat_nul_as_a_byte(self, dynamodb_client, nul_table): + below_space = _query( + dynamodb_client, nul_table, cond="pk = :p AND sk < :s", values={":s": {"S": " "}}, + ConsistentRead=True, + ) + assert _sks(below_space) == [NUL, SOH] + between = _query( + dynamodb_client, nul_table, cond="pk = :p AND sk BETWEEN :lo AND :hi", + values={":lo": {"S": NUL}, ":hi": {"S": SOH}}, ConsistentRead=True, + ) + assert _sks(between) == [NUL, SOH] + prefix = _query( + dynamodb_client, nul_table, cond="pk = :p AND begins_with(sk, :pre)", + values={":pre": {"S": NUL}}, ConsistentRead=True, + ) + assert _sks(prefix) == [NUL] + text_prefix = _query( + dynamodb_client, nul_table, cond="pk = :p AND begins_with(sk, :pre)", + values={":pre": {"S": "\\"}}, ConsistentRead=True, + ) + assert _sks(text_prefix) == [ESCAPE_TEXT] + + def test_pages_of_one_visit_every_key_once(self, dynamodb_client, nul_table): + """The cursor for a NUL key travels through LastEvaluatedKey and back + through ExclusiveStartKey; the walk must neither skip nor repeat.""" + seen: list[str] = [] + kwargs: dict = {"Limit": 1, "ConsistentRead": True} + while True: + resp = _query(dynamodb_client, nul_table, **kwargs) + seen += _sks(resp) + if "LastEvaluatedKey" not in resp: + break + kwargs["ExclusiveStartKey"] = resp["LastEvaluatedKey"] + assert seen == SORT_KEYS + + def test_gsi_holds_the_item_with_a_nul_key(self, dynamodb_client, nul_table): + def walk(): + return dynamodb_client.query( + TableName=nul_table, + IndexName="gsi1", + KeyConditionExpression="g = :g", + ExpressionAttributeValues={":g": {"S": f"g{NUL}"}}, + )["Items"] + + items = wait_for_gsi_items(walk, 1) + assert items == [RICH_ITEM] + + def test_filter_on_a_nul_named_attribute(self, dynamodb_client, nul_table): + resp = dynamodb_client.scan( + TableName=nul_table, + FilterExpression="#n = :x", + ExpressionAttributeNames={"#n": f"n{NUL}n"}, + ExpressionAttributeValues={":x": {"S": "x"}}, + ConsistentRead=True, + ) + assert resp["Items"] == [RICH_ITEM] + resp = _query( + dynamodb_client, nul_table, FilterExpression="v = :v", + values={":v": {"S": f"v{NUL}"}}, ConsistentRead=True, + ) + assert resp["Items"] == [RICH_ITEM] + + def test_condition_expression_sees_the_nul_named_attribute(self, dynamodb_client, nul_table): + # Exists on the rich item: the conditional write goes through. + dynamodb_client.update_item( + TableName=nul_table, + Key={"pk": {"S": PK}, "sk": {"S": NUL}}, + UpdateExpression="SET touched = :t", + ConditionExpression="attribute_exists(#n)", + ExpressionAttributeNames={"#n": f"n{NUL}n"}, + ExpressionAttributeValues={":t": {"BOOL": True}}, + ) + with pytest.raises(ClientError) as exc: + dynamodb_client.update_item( + TableName=nul_table, + Key={"pk": {"S": PK}, "sk": {"S": "a"}}, + UpdateExpression="SET touched = :t", + ConditionExpression="attribute_exists(#n)", + ExpressionAttributeNames={"#n": f"n{NUL}n"}, + ExpressionAttributeValues={":t": {"BOOL": True}}, + ) + assert exc.value.response["Error"]["Code"] == "ConditionalCheckFailedException" + dynamodb_client.update_item( + TableName=nul_table, + Key={"pk": {"S": PK}, "sk": {"S": NUL}}, + UpdateExpression="REMOVE touched", + ) + + def test_update_item_sets_nul_named_attribute_and_nul_value(self, dynamodb_client, nul_table): + key = {"pk": {"S": PK}, "sk": {"S": " "}} + resp = dynamodb_client.update_item( + TableName=nul_table, + Key=key, + UpdateExpression="SET #a = :v, m2 = :m", + ExpressionAttributeNames={"#a": f"u{NUL}"}, + ExpressionAttributeValues={ + ":v": {"S": f"{NUL}z"}, + ":m": {"M": {f"{NUL}": {"L": [{"S": SOH}]}}}, + }, + ReturnValues="ALL_NEW", + ) + assert resp["Attributes"][f"u{NUL}"] == {"S": f"{NUL}z"} + assert resp["Attributes"]["m2"] == {"M": {f"{NUL}": {"L": [{"S": SOH}]}}} + got = dynamodb_client.get_item(TableName=nul_table, Key=key, ConsistentRead=True)["Item"] + assert got[f"u{NUL}"] == {"S": f"{NUL}z"} + assert got["m2"] == {"M": {f"{NUL}": {"L": [{"S": SOH}]}}} + resp = dynamodb_client.update_item( + TableName=nul_table, + Key=key, + UpdateExpression="REMOVE #a, m2", + ExpressionAttributeNames={"#a": f"u{NUL}"}, + ReturnValues="ALL_NEW", + ) + assert f"u{NUL}" not in resp["Attributes"] and "m2" not in resp["Attributes"] + + def test_batch_and_transact_writes_and_reads(self, dynamodb_client, nul_table): + b1 = {"pk": {"S": f"batch{NUL}"}, "sk": {"S": NUL}, "g": {"S": "plain"}, f"{NUL}": {"S": NUL}} + t1 = {"pk": {"S": f"tx{NUL}"}, "sk": {"S": f"{NUL}{NUL}"}, "g": {"S": "plain"}} + dynamodb_client.batch_write_item(RequestItems={nul_table: [{"PutRequest": {"Item": b1}}]}) + dynamodb_client.transact_write_items( + TransactItems=[{"Put": {"TableName": nul_table, "Item": t1}}] + ) + got = dynamodb_client.batch_get_item( + RequestItems={ + nul_table: { + "Keys": [{"pk": b1["pk"], "sk": b1["sk"]}, {"pk": t1["pk"], "sk": t1["sk"]}], + "ConsistentRead": True, + } + } + )["Responses"][nul_table] + assert sorted(got, key=lambda i: i["pk"]["S"]) == sorted([b1, t1], key=lambda i: i["pk"]["S"]) + tg = dynamodb_client.transact_get_items( + TransactItems=[ + {"Get": {"TableName": nul_table, "Key": {"pk": b1["pk"], "sk": b1["sk"]}}}, + {"Get": {"TableName": nul_table, "Key": {"pk": t1["pk"], "sk": t1["sk"]}}}, + ] + )["Responses"] + assert [r["Item"] for r in tg] == [b1, t1] + dynamodb_client.batch_write_item( + RequestItems={ + nul_table: [ + {"DeleteRequest": {"Key": {"pk": b1["pk"], "sk": b1["sk"]}}}, + {"DeleteRequest": {"Key": {"pk": t1["pk"], "sk": t1["sk"]}}}, + ] + } + ) + + def test_scan_counts_every_item(self, dynamodb_client, nul_table): + resp = dynamodb_client.scan(TableName=nul_table, Select="COUNT", ConsistentRead=True) + assert resp["Count"] == len(SORT_KEYS) + 2 + + def test_delete_item_on_a_nul_key(self, dynamodb_client, nul_table): + key = {"pk": {"S": f"del{NUL}"}, "sk": {"S": NUL}} + dynamodb_client.put_item(TableName=nul_table, Item={**key, "g": {"S": "plain"}}) + resp = dynamodb_client.delete_item(TableName=nul_table, Key=key, ReturnValues="ALL_OLD") + assert resp["Attributes"] == {**key, "g": {"S": "plain"}} + assert "Item" not in dynamodb_client.get_item(TableName=nul_table, Key=key, ConsistentRead=True) + + +@pytest.fixture(scope="module") +def nul_streams_client(endpoint_url): + """Streams client on the ExtendDB endpoint. The service routes Streams to a + separate endpoint boto3 resolves on its own, so this test is ExtendDB only, + as the stream tests in ``test_streams.py`` are.""" + if not endpoint_url: + pytest.skip("stream record fidelity is checked against ExtendDB servers only") + import os + + import boto3 + + kwargs: dict = dict( + service_name="dynamodbstreams", + region_name=os.environ.get("AWS_DEFAULT_REGION", "us-east-1"), + endpoint_url=endpoint_url, + ) + if endpoint_url.startswith("https://"): + kwargs["verify"] = False + return boto3.client(**kwargs) + + +def test_stream_records_carry_nul_keys_names_and_values(dynamodb_client, nul_streams_client): + """INSERT, MODIFY, and REMOVE records for an item whose key, attribute + names, map keys, and values hold U+0000 and U+0001 come back byte-identical + in Keys, NewImage, and OldImage.""" + from test_streams import _drain_all_shards + + with scoped_table( + dynamodb_client, + attribute_definitions=[ + {"AttributeName": "pk", "AttributeType": "S"}, + {"AttributeName": "sk", "AttributeType": "S"}, + ], + key_schema=[ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "sk", "KeyType": "RANGE"}, + ], + StreamSpecification={"StreamEnabled": True, "StreamViewType": "NEW_AND_OLD_IMAGES"}, + ) as name: + stream_arn = dynamodb_client.describe_table(TableName=name)["Table"]["LatestStreamArn"] + key = {"pk": {"S": PK}, "sk": {"S": NUL}} + first = {**RICH_ITEM} + dynamodb_client.put_item(TableName=name, Item=first) + second = dynamodb_client.update_item( + TableName=name, + Key=key, + UpdateExpression="SET #a = :v REMOVE #n", + ExpressionAttributeNames={"#a": f"u{NUL}", "#n": f"n{NUL}n"}, + ExpressionAttributeValues={":v": {"S": f"{SOH}{NUL}"}}, + ReturnValues="ALL_NEW", + )["Attributes"] + dynamodb_client.delete_item(TableName=name, Key=key) + + deadline = time.monotonic() + 30 + records: list[dict] = [] + while time.monotonic() < deadline: + records = _drain_all_shards(nul_streams_client, stream_arn) + if len(records) >= 3: + break + time.sleep(0.5) + assert [r["eventName"] for r in records] == ["INSERT", "MODIFY", "REMOVE"] + for r in records: + assert r["dynamodb"]["Keys"] == key + assert records[0]["dynamodb"]["NewImage"] == first + assert "OldImage" not in records[0]["dynamodb"] + assert records[1]["dynamodb"]["OldImage"] == first + assert records[1]["dynamodb"]["NewImage"] == second + assert records[2]["dynamodb"]["OldImage"] == second + assert "NewImage" not in records[2]["dynamodb"]