Skip to content

Commit 5f64ddd

Browse files
farhan-syahfarhan-syah
committed
security(sql): make procedural binding UTF-8 safe
1 parent 923fdec commit 5f64ddd

5 files changed

Lines changed: 667 additions & 199 deletions

File tree

‎nodedb/src/control/planner/procedural/compiler.rs‎

Lines changed: 138 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -46,12 +46,7 @@ impl CompileContext {
4646

4747
/// Substitute known variables in a SQL expression.
4848
fn substitute(&self, sql: &str) -> String {
49-
let mut result = sql.to_string();
50-
for (name, expr) in &self.variables {
51-
// Replace whole-word occurrences of the variable name.
52-
result = replace_identifier(&result, name, &format!("({expr})"));
53-
}
54-
result
49+
substitute_identifiers(sql, &self.variables)
5550
}
5651
}
5752

@@ -82,11 +77,11 @@ fn compile_statements(
8277
Some(e) => ctx.substitute(&e.sql),
8378
None => "NULL".into(),
8479
};
85-
ctx.variables.insert(name.clone(), expr);
80+
ctx.variables.insert(name.to_lowercase(), expr);
8681
}
8782
Statement::Assign { target, expr } => {
8883
let substituted = ctx.substitute(&expr.sql);
89-
ctx.variables.insert(target.clone(), substituted);
84+
ctx.variables.insert(target.to_lowercase(), substituted);
9085
}
9186
Statement::If {
9287
condition,
@@ -240,7 +235,9 @@ fn compile_for(
240235
let mut case_parts = Vec::new();
241236
for val in &iterations {
242237
let mut iter_ctx = ctx.clone_vars();
243-
iter_ctx.variables.insert(var.to_string(), val.to_string());
238+
iter_ctx
239+
.variables
240+
.insert(var.to_lowercase(), val.to_string());
244241
let body_sql = compile_statements(body, &mut iter_ctx)?;
245242
// Each iteration becomes a WHEN clause. If the body is unconditional
246243
// (always returns), it's `WHEN true THEN <body>`. If conditional,
@@ -272,70 +269,125 @@ impl CompileContext {
272269
}
273270
}
274271

275-
/// Replace whole-word occurrences of `name` in `sql` with `replacement`.
276-
///
277-
/// Only replaces when `name` is surrounded by non-alphanumeric characters
278-
/// (or at string boundaries). This prevents replacing "x" in "max" or "text".
279-
fn replace_identifier(sql: &str, name: &str, replacement: &str) -> String {
280-
if name.is_empty() || !sql.contains(name) {
281-
return sql.to_string();
282-
}
283-
284-
let mut result = String::with_capacity(sql.len());
285-
let bytes = sql.as_bytes();
286-
let name_bytes = name.as_bytes();
287-
let name_len = name_bytes.len();
288-
let mut i = 0;
289-
290-
while i < bytes.len() {
291-
// Skip string literals.
292-
if bytes[i] == b'\'' {
293-
result.push('\'');
294-
i += 1;
295-
while i < bytes.len() {
296-
if bytes[i] == b'\'' {
297-
result.push('\'');
298-
i += 1;
299-
if i < bytes.len() && bytes[i] == b'\'' {
300-
result.push('\'');
301-
i += 1;
302-
} else {
303-
break;
304-
}
272+
/// Substitute all known variables in one lexical pass. Replacement text is
273+
/// copied, never revisited, so binding expansion is deterministic.
274+
fn substitute_identifiers(sql: &str, variables: &HashMap<String, String>) -> String {
275+
let mut output = String::with_capacity(sql.len());
276+
let mut cursor = 0;
277+
while cursor < sql.len() {
278+
let end = opaque_end(sql, cursor);
279+
if end > cursor {
280+
output.push_str(&sql[cursor..end]);
281+
cursor = end;
282+
} else if let Some(ch) = sql[cursor..].chars().next() {
283+
if is_identifier_start(ch) {
284+
let end = consume_identifier(sql, cursor);
285+
let token = &sql[cursor..end];
286+
if let Some(value) = variables.get(&token.to_lowercase()) {
287+
output.push('(');
288+
output.push_str(value);
289+
output.push(')');
305290
} else {
306-
result.push(bytes[i] as char);
307-
i += 1;
291+
output.push_str(token);
308292
}
293+
cursor = end;
294+
} else {
295+
output.push(ch);
296+
cursor += ch.len_utf8();
309297
}
310-
continue;
311298
}
299+
}
300+
output
301+
}
302+
303+
fn opaque_end(sql: &str, start: usize) -> usize {
304+
let bytes = sql.as_bytes();
305+
match bytes[start] {
306+
b'\'' | b'"' => consume_quoted(sql, start, bytes[start]),
307+
b'-' if bytes.get(start + 1) == Some(&b'-') => sql[start..]
308+
.find('\n')
309+
.map_or(sql.len(), |offset| start + offset),
310+
b'/' if bytes.get(start + 1) == Some(&b'*') => consume_block_comment(sql, start),
311+
b'$' => consume_dollar_quote(sql, start).unwrap_or(start),
312+
_ => start,
313+
}
314+
}
312315

313-
// Check for whole-word match.
314-
if i + name_len <= bytes.len()
315-
&& sql[i..i + name_len].eq_ignore_ascii_case(name)
316-
&& !is_ident_char(bytes.get(i.wrapping_sub(1)).copied(), i == 0)
317-
&& !is_ident_continue(bytes.get(i + name_len).copied())
318-
{
319-
result.push_str(replacement);
320-
i += name_len;
316+
fn consume_quoted(sql: &str, start: usize, quote: u8) -> usize {
317+
let bytes = sql.as_bytes();
318+
let mut cursor = start + 1;
319+
while cursor < bytes.len() {
320+
if bytes[cursor] == quote {
321+
cursor += 1;
322+
if bytes.get(cursor) != Some(&quote) {
323+
return cursor;
324+
}
325+
cursor += 1;
321326
} else {
322-
result.push(bytes[i] as char);
323-
i += 1;
327+
cursor += sql[cursor..].chars().next().map_or(1, char::len_utf8);
324328
}
325329
}
330+
sql.len()
331+
}
326332

327-
result
333+
fn consume_block_comment(sql: &str, start: usize) -> usize {
334+
let bytes = sql.as_bytes();
335+
let mut cursor = start + 2;
336+
let mut depth = 1usize;
337+
while cursor + 1 < bytes.len() {
338+
match (bytes[cursor], bytes[cursor + 1]) {
339+
(b'/', b'*') => {
340+
depth += 1;
341+
cursor += 2;
342+
}
343+
(b'*', b'/') => {
344+
depth -= 1;
345+
cursor += 2;
346+
if depth == 0 {
347+
return cursor;
348+
}
349+
}
350+
_ => cursor += sql[cursor..].chars().next().map_or(1, char::len_utf8),
351+
}
352+
}
353+
sql.len()
328354
}
329355

330-
fn is_ident_char(byte: Option<u8>, at_start: bool) -> bool {
331-
if at_start {
332-
return false;
356+
fn consume_dollar_quote(sql: &str, start: usize) -> Option<usize> {
357+
let rest = &sql[start..];
358+
let close = rest[1..].find('$')? + 1;
359+
let tag = &rest[..=close];
360+
let tag_body = &tag[1..tag.len() - 1];
361+
if !tag_body.is_empty()
362+
&& (!tag_body
363+
.chars()
364+
.next()
365+
.is_some_and(|ch| ch == '_' || ch.is_alphabetic())
366+
|| !tag_body.chars().all(|ch| ch == '_' || ch.is_alphanumeric()))
367+
{
368+
return None;
333369
}
334-
byte.is_some_and(|b| b.is_ascii_alphanumeric() || b == b'_')
370+
let body_start = start + tag.len();
371+
Some(
372+
sql[body_start..]
373+
.find(tag)
374+
.map_or(sql.len(), |offset| body_start + offset + tag.len()),
375+
)
335376
}
336377

337-
fn is_ident_continue(byte: Option<u8>) -> bool {
338-
byte.is_some_and(|b| b.is_ascii_alphanumeric() || b == b'_')
378+
fn is_identifier_start(ch: char) -> bool {
379+
ch == '_' || ch.is_alphabetic()
380+
}
381+
fn is_identifier_continue(ch: char) -> bool {
382+
ch == '_' || ch == '$' || ch.is_alphanumeric()
383+
}
384+
fn consume_identifier(sql: &str, start: usize) -> usize {
385+
sql[start..]
386+
.char_indices()
387+
.find_map(|(offset, ch)| {
388+
(offset > 0 && !is_identifier_continue(ch)).then_some(start + offset)
389+
})
390+
.unwrap_or(sql.len())
339391
}
340392

341393
#[cfg(test)]
@@ -472,11 +524,33 @@ mod tests {
472524
}
473525

474526
#[test]
475-
fn replace_whole_word() {
476-
assert_eq!(replace_identifier("x + 1", "x", "(5)"), "(5) + 1");
477-
assert_eq!(replace_identifier("max(x)", "x", "(5)"), "max((5))");
478-
// Should NOT replace "x" inside "text" or "max".
479-
assert_eq!(replace_identifier("text", "x", "(5)"), "text");
527+
fn replaces_only_whole_identifier_tokens() {
528+
let variables = HashMap::from([("x".to_string(), "5".to_string())]);
529+
assert_eq!(substitute_identifiers("x + 1", &variables), "(5) + 1");
530+
assert_eq!(substitute_identifiers("max(x)", &variables), "max((5))");
531+
assert_eq!(substitute_identifiers("text", &variables), "text");
532+
}
533+
534+
#[test]
535+
fn replacement_is_one_pass_and_respects_unicode_dollar_boundaries() {
536+
let mut variables = HashMap::new();
537+
variables.insert("x".into(), "y".into());
538+
variables.insert("y".into(), "9".into());
539+
variables.insert("é".into(), "11".into());
540+
assert_eq!(
541+
substitute_identifiers("x y αx xβ x$ É é", &variables),
542+
"(y) (9) αx xβ x$ (11) (11)"
543+
);
544+
}
545+
546+
#[test]
547+
fn replacement_preserves_unicode_and_opaque_sql_regions() {
548+
let variables = HashMap::from([("x".to_string(), "5".to_string())]);
549+
let sql = "x + '雪x' + \"x\" -- x 雪\n/* outer x /* inner x */ 雪 */ + x";
550+
assert_eq!(
551+
substitute_identifiers(sql, &variables),
552+
"(5) + '雪x' + \"x\" -- x 雪\n/* outer x /* inner x */ 雪 */ + (5)"
553+
);
480554
}
481555

482556
#[test]

0 commit comments

Comments
 (0)