Skip to content

Commit 1341b39

Browse files
V02460abonander
andauthored
Escape PostgreSQL options (#3800)
* Escape PostgreSQL options * Use raw string literals in test case Co-authored-by: Austin Bonander <austin.bonander@gmail.com> * Document escaping behavior for options * Remove heap allocations for options formatting * Use an actual config option for the test case --------- Co-authored-by: Austin Bonander <austin.bonander@gmail.com>
1 parent 25cbeed commit 1341b39

File tree

1 file changed

+59
-2
lines changed
  • sqlx-postgres/src/options

1 file changed

+59
-2
lines changed

sqlx-postgres/src/options/mod.rs

Lines changed: 59 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use std::borrow::Cow;
22
use std::env::var;
3-
use std::fmt::{Display, Write};
3+
use std::fmt::{self, Display, Write};
44
use std::path::{Path, PathBuf};
55

66
pub use ssl_mode::PgSslMode;
@@ -416,6 +416,9 @@ impl PgConnectOptions {
416416

417417
/// Set additional startup options for the connection as a list of key-value pairs.
418418
///
419+
/// Escapes the options’ backslash and space characters as per
420+
/// https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNECT-OPTIONS
421+
///
419422
/// # Example
420423
///
421424
/// ```rust
@@ -436,7 +439,8 @@ impl PgConnectOptions {
436439
options_str.push(' ');
437440
}
438441

439-
write!(options_str, "-c {k}={v}").expect("failed to write an option to the string");
442+
options_str.push_str("-c ");
443+
write!(PgOptionsWriteEscaped(options_str), "{k}={v}").ok();
440444
}
441445
self
442446
}
@@ -590,6 +594,39 @@ fn default_host(port: u16) -> String {
590594
"localhost".to_owned()
591595
}
592596

597+
/// Writer that escapes passed-in PostgreSQL options.
598+
///
599+
/// Escapes backslashes and spaces with an additional backslash according to
600+
/// https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNECT-OPTIONS
601+
#[derive(Debug)]
602+
struct PgOptionsWriteEscaped<'a>(&'a mut String);
603+
604+
impl Write for PgOptionsWriteEscaped<'_> {
605+
fn write_str(&mut self, s: &str) -> fmt::Result {
606+
let mut span_start = 0;
607+
608+
for (span_end, matched) in s.match_indices([' ', '\\']) {
609+
write!(self.0, r"{}\{matched}", &s[span_start..span_end])?;
610+
span_start = span_end + matched.len();
611+
}
612+
613+
// Write the rest of the string after the last match, or all of it if no matches
614+
self.0.push_str(&s[span_start..]);
615+
616+
Ok(())
617+
}
618+
619+
fn write_char(&mut self, ch: char) -> fmt::Result {
620+
if matches!(ch, ' ' | '\\') {
621+
self.0.push('\\');
622+
}
623+
624+
self.0.push(ch);
625+
626+
Ok(())
627+
}
628+
}
629+
593630
#[test]
594631
fn test_options_formatting() {
595632
let options = PgConnectOptions::new().options([("geqo", "off")]);
@@ -604,6 +641,26 @@ fn test_options_formatting() {
604641
options.options,
605642
Some("-c geqo=off -c statement_timeout=5min".to_string())
606643
);
644+
// https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNECT-OPTIONS
645+
let options =
646+
PgConnectOptions::new().options([("application_name", r"/back\slash/ and\ spaces")]);
647+
assert_eq!(
648+
options.options,
649+
Some(r"-c application_name=/back\\slash/\ and\\\ spaces".to_string())
650+
);
607651
let options = PgConnectOptions::new();
608652
assert_eq!(options.options, None);
609653
}
654+
655+
#[test]
656+
fn test_pg_write_escaped() {
657+
let mut buf = String::new();
658+
let mut x = PgOptionsWriteEscaped(&mut buf);
659+
x.write_str("x").unwrap();
660+
x.write_str("").unwrap();
661+
x.write_char('\\').unwrap();
662+
x.write_str("y \\").unwrap();
663+
x.write_char(' ').unwrap();
664+
x.write_char('z').unwrap();
665+
assert_eq!(buf, r"x\\y\ \\\ z");
666+
}

0 commit comments

Comments
 (0)