Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit c5ae79a

Browse files
authored
Merge pull request sfackler#702 from jeff-davis/escape
Add escape_identifier() and escape_literal().
2 parents 83ffcf8 + aadf942 commit c5ae79a

File tree

3 files changed

+111
-0
lines changed

3 files changed

+111
-0
lines changed

postgres-protocol/src/escape/mod.rs

+93
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
//! Provides functions for escaping literals and identifiers for use
2+
//! in SQL queries.
3+
//!
4+
//! Prefer parameterized queries where possible; see
5+
//! [`Client::query`](crate::Client::query). Do not escape parameters.
6+
7+
#[cfg(test)]
8+
mod test;
9+
10+
/// Escape a literal and surround result with single quotes. Not
11+
/// recommended in most cases.
12+
///
13+
/// If input contains backslashes, result will be of the form `
14+
/// E'...'` so it is safe to use regardless of the setting of
15+
/// standard_conforming_strings.
16+
pub fn escape_literal(input: &str) -> String {
17+
escape_internal(input, false)
18+
}
19+
20+
/// Escape an identifier and surround result with double quotes.
21+
pub fn escape_identifier(input: &str) -> String {
22+
escape_internal(input, true)
23+
}
24+
25+
// Translation of PostgreSQL libpq's PQescapeInternal(). Does not
26+
// require a connection because input string is known to be valid
27+
// UTF-8.
28+
//
29+
// Escape arbitrary strings. If as_ident is true, we escape the
30+
// result as an identifier; if false, as a literal. The result is
31+
// returned in a newly allocated buffer. If we fail due to an
32+
// encoding violation or out of memory condition, we return NULL,
33+
// storing an error message into conn.
34+
fn escape_internal(input: &str, as_ident: bool) -> String {
35+
let mut num_backslashes = 0;
36+
let mut num_quotes = 0;
37+
let quote_char = if as_ident { '"' } else { '\'' };
38+
39+
// Scan the string for characters that must be escaped.
40+
for ch in input.chars() {
41+
if ch == quote_char {
42+
num_quotes += 1;
43+
} else if ch == '\\' {
44+
num_backslashes += 1;
45+
}
46+
}
47+
48+
// Allocate output String.
49+
let mut result_size = input.len() + num_quotes + 3; // two quotes, plus a NUL
50+
if !as_ident && num_backslashes > 0 {
51+
result_size += num_backslashes + 2;
52+
}
53+
54+
let mut output = String::with_capacity(result_size);
55+
56+
// If we are escaping a literal that contains backslashes, we use
57+
// the escape string syntax so that the result is correct under
58+
// either value of standard_conforming_strings. We also emit a
59+
// leading space in this case, to guard against the possibility
60+
// that the result might be interpolated immediately following an
61+
// identifier.
62+
if !as_ident && num_backslashes > 0 {
63+
output.push(' ');
64+
output.push('E');
65+
}
66+
67+
// Opening quote.
68+
output.push(quote_char);
69+
70+
// Use fast path if possible.
71+
//
72+
// We've already verified that the input string is well-formed in
73+
// the current encoding. If it contains no quotes and, in the
74+
// case of literal-escaping, no backslashes, then we can just copy
75+
// it directly to the output buffer, adding the necessary quotes.
76+
//
77+
// If not, we must rescan the input and process each character
78+
// individually.
79+
if num_quotes == 0 && (num_backslashes == 0 || as_ident) {
80+
output.push_str(input);
81+
} else {
82+
for ch in input.chars() {
83+
if ch == quote_char || (!as_ident && ch == '\\') {
84+
output.push(ch);
85+
}
86+
output.push(ch);
87+
}
88+
}
89+
90+
output.push(quote_char);
91+
92+
output
93+
}

postgres-protocol/src/escape/test.rs

+17
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
use crate::escape::{escape_identifier, escape_literal};
2+
3+
#[test]
4+
fn test_escape_idenifier() {
5+
assert_eq!(escape_identifier("foo"), String::from("\"foo\""));
6+
assert_eq!(escape_identifier("f\\oo"), String::from("\"f\\oo\""));
7+
assert_eq!(escape_identifier("f'oo"), String::from("\"f'oo\""));
8+
assert_eq!(escape_identifier("f\"oo"), String::from("\"f\"\"oo\""));
9+
}
10+
11+
#[test]
12+
fn test_escape_literal() {
13+
assert_eq!(escape_literal("foo"), String::from("'foo'"));
14+
assert_eq!(escape_literal("f\\oo"), String::from(" E'f\\\\oo'"));
15+
assert_eq!(escape_literal("f'oo"), String::from("'f''oo'"));
16+
assert_eq!(escape_literal("f\"oo"), String::from("'f\"oo'"));
17+
}

postgres-protocol/src/lib.rs

+1
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ use bytes::{BufMut, BytesMut};
1717
use std::io;
1818

1919
pub mod authentication;
20+
pub mod escape;
2021
pub mod message;
2122
pub mod types;
2223

0 commit comments

Comments
 (0)