Properly escape table and column names in prepare_copy_in

We have to assemble queries by hand here which is a bit sketchy.
Manually escaping the individual identifiers to avoid introducing
injection vulernabilities is unfortunate but necessary.

Author: sfackler

src/lib.rs

@@ -763,15 +763,18 @@ impl InnerConnection {
                            -> Result<CopyInStatement<'a>> {
         let mut query = vec![];
         let _ = write!(&mut query, "SELECT ");
-        let _ = util::comma_join(&mut query, rows.iter().cloned());
-        let _ = write!(&mut query, " FROM {}", table);
+        let _ = util::comma_join_quoted_idents(&mut query, rows.iter().cloned());
+        let _ = write!(&mut query, " FROM ");
+        let _ = util::write_quoted_ident(&mut query, table);
         let query = String::from_utf8(query).unwrap();
         let (_, columns) = try!(self.raw_prepare("", &query));
         let column_types = columns.into_iter().map(|desc| desc.type_).collect();
 
         let mut query = vec![];
-        let _ = write!(&mut query, "COPY {} (", table);
-        let _ = util::comma_join(&mut query, rows.iter().cloned());
+        let _ = write!(&mut query, "COPY ");
+        let _ = util::write_quoted_ident(&mut query, table);
+        let _ = write!(&mut query, " (");
+        let _ = util::comma_join_quoted_idents(&mut query, rows.iter().cloned());
         let _ = write!(&mut query, ") FROM STDIN WITH (FORMAT binary)");
         let query = String::from_utf8(query).unwrap();
         let stmt_name = self.make_stmt_name();

src/util.rs

@@ -1,19 +1,34 @@
 use std::io;
 use std::io::prelude::*;
+use std::ascii::AsciiExt;
 
-pub fn comma_join<'a, W, I>(writer: &mut W, strs: I) -> io::Result<()>
+pub fn comma_join_quoted_idents<'a, W, I>(writer: &mut W, strs: I) -> io::Result<()>
         where W: Write, I: Iterator<Item=&'a str> {
     let mut first = true;
     for str_ in strs {
         if !first {
             try!(write!(writer, ", "));
         }
         first = false;
-        try!(write!(writer, "{}", str_));
+        try!(write_quoted_ident(writer, str_));
     }
     Ok(())
 }
 
+// See http://www.postgresql.org/docs/9.4/static/sql-syntax-lexical.html for ident grammar
+pub fn write_quoted_ident<W: Write>(w: &mut W, ident: &str) -> io::Result<()> {
+    try!(write!(w, "U&\""));
+    for ch in ident.chars() {
+        match ch {
+            '"' => try!(write!(w, "\"\"")),
+            '\\' => try!(write!(w, "\\\\")),
+            ch if ch.is_ascii() => try!(write!(w, "{}", ch)),
+            ch => try!(write!(w, "\\+{:06X}", ch as u32)),
+        }
+    }
+    write!(w, "\"")
+}
+
 pub fn parse_update_count(tag: String) -> u64 {
     tag.split(' ').last().unwrap().parse().unwrap_or(0)
 }

tests/test.rs

@@ -32,7 +32,7 @@ macro_rules! or_panic {
     ($e:expr) => (
         match $e {
             Ok(ok) => ok,
-            Err(err) => panic!("{:?}", err)
+            Err(err) => panic!("{:#?}", err)
         }
     )
 }
@@ -832,6 +832,14 @@ fn test_copy_in_bad_type() {
     or_panic!(conn.execute("SELECT 1", &[]));
 }
 
+#[test]
+fn test_copy_in_weird_names() {
+    let conn = or_panic!(Connection::connect("postgres://postgres@localhost", &SslMode::None));
+    or_panic!(conn.execute(r#"CREATE TEMPORARY TABLE "na""me" (U&" \\\+01F4A9" VARCHAR)"#, &[]));
+    let stmt = or_panic!(conn.prepare_copy_in("na\"me", &[" \\💩"]));
+    assert_eq!(&Type::Varchar, &stmt.column_types()[0]);
+}
+
 #[test]
 fn test_batch_execute_copy_from_err() {
     let conn = or_panic!(Connection::connect("postgres://postgres@localhost", &SslMode::None));